Tracing

cfg = L.call(Model)(encoder=L.call(Encoder)(depth=24, width=512), dim=768)

cfg = L.trace(lambda: Model(encoder=Encoder(depth=24, width=512), dim=768))

import laco
import laco.language as L
from omegaconf import DictConfig

# --- Making toy classes importable -----------------------------------------
# laco records an *importable* target path (module.Qualname) for every node so
# the config can be dumped to YAML and re-instantiated later. A class defined
# in a notebook cell lives in the ``__main__`` module and is NOT importable,
# so ``laco.dump`` would fail with "Cannot generate path for object ...".
#
# To keep the toy classes below readable while still teaching the real
# round-trip, we register each one into a small, genuinely-importable module.
# The ``@importable`` decorator below does exactly that: it rehomes a
# notebook-defined class into ``nb_tracing_demo`` so ``module.Qualname`` resolves.
import sys
import types

_demo = types.ModuleType("nb_tracing_demo")
sys.modules["nb_tracing_demo"] = _demo


def importable(cls):
    """Rehome a notebook class into an importable module so laco can dump it."""
    cls.__module__ = _demo.__name__
    setattr(_demo, cls.__qualname__, cls)
    return cls

# Define toy classes (no torch, no heavy deps).
# @importable rehomes them into the nb_tracing_demo module so laco can record
# an importable target path (module.Qualname) — see the imports cell.
@importable
class Encoder:
    def __init__(self, depth: int, width: int = 512):
        self.depth = depth
        self.width = width

@importable
class Model:
    def __init__(self, encoder: Encoder, dim: int = 512):
        self.encoder = encoder
        self.dim = dim

# --- Approach 1: Explicit L.call ---
cfg_explicit = L.call(Model)(
    encoder=L.call(Encoder)(depth=24, width=512),
    dim=768
)

print("=== Approach 1: explicit L.call ===")
print(laco.dump(cfg_explicit))

# --- Approach 2: Tracing ---
# Decorate the classes with @L.configurable
TraceEncoder = L.configurable(Encoder)
TraceModel   = L.configurable(Model)

cfg_traced = L.trace(lambda: TraceModel(
    encoder=TraceEncoder(depth=24, width=512),
    dim=768
))

print("=== Approach 2: L.trace ===")
print(laco.dump(cfg_traced))

# Both produce the same YAML:
assert laco.dump(cfg_explicit) == laco.dump(cfg_traced), "Should be identical!"
print("Both approaches produce identical YAML.")

# Both can be instantiated the same way:
model_from_explicit = laco.instantiate(cfg_explicit)
model_from_trace    = laco.instantiate(cfg_traced)
print(f"Instantiated from explicit: {type(model_from_explicit).__name__}, dim={model_from_explicit.dim}")  # noqa: LACO001
print(f"Instantiated from trace   : {type(model_from_trace).__name__}, dim={model_from_trace.dim}")        # noqa: LACO001

@L.configurable
@importable
class SimpleEncoder:
    def __init__(self, depth: int, width: int = 512):
        self.depth = depth
        self.width = width
    def __repr__(self):
        return f"SimpleEncoder(depth={self.depth}, width={self.width})"

# --- OUTSIDE a trace: executes normally, returns a real object ---
enc_outside = SimpleEncoder(depth=8)
print("Outside trace:")
print(f"  type : {type(enc_outside).__name__}")
print(f"  repr : {enc_outside}")
print(f"  depth: {enc_outside.depth}")  # noqa: LACO001

# --- INSIDE a trace: intercepts the call, returns a DictConfig ---
cfg_inside = L.trace(lambda: SimpleEncoder(depth=8))
print("Inside trace:")
print(f"  type    : {type(cfg_inside).__name__}")
print(f"  depth   : {cfg_inside.depth}")
# noqa: LACO001  -- intentional demo of config key access
print(f"  width   : {cfg_inside.width}")
# noqa: LACO001
print(f"  _target_: {cfg_inside._target_}")
# noqa: LACO001
print()
print(laco.dump(cfg_inside))

# Factory form: @L.configurable(strict=False) — for targets with unintrospectable signatures
@L.configurable(strict=False)
class FlexibleLayer:
    def __init__(self, units: int, activation: str = "relu"):
        self.units = units
        self.activation = activation

# Outside trace — works normally:
layer_live = FlexibleLayer(units=64)
print(f"Outside trace: {type(layer_live).__name__}, units={layer_live.units}")  # noqa: LACO001

# Inside trace — records a node:
cfg_layer = L.trace(lambda: FlexibleLayer(units=64, activation="gelu"))
print(f"Inside trace: {type(cfg_layer).__name__}, units={cfg_layer.units}")  # noqa: LACO001

from laco.language import _TRACING

print("_TRACING type   :", type(_TRACING).__name__)
print("default (outside):", _TRACING.get())  # False outside any trace

# Illustrate the token-based set/reset mechanism:
print("Before trace:", _TRACING.get())

# What L.trace(thunk) does internally:
def demo_trace_mechanism(thunk):
    token = _TRACING.set(True)   # set to True, save token
    try:
        print("  Inside trace:", _TRACING.get())
        result = thunk()
        return result
    finally:
        _TRACING.reset(token)    # always reset, even on exception
        print("  After finally:", _TRACING.get())

demo_trace_mechanism(lambda: None)
print("After trace:", _TRACING.get())

# The ContextVar resets even when the thunk raises:
print("Before exception trace:", _TRACING.get())

try:
    L.trace(lambda: (_ for _ in ()).throw(RuntimeError("boom")))
except (RuntimeError, TypeError):
    pass  # expected

print("After exception trace:", _TRACING.get())  # still False — no leak!

@L.configurable
@importable
class TieredContainer:
    def __init__(self, node_arg, int_arg: int, str_arg: str):
        pass

# --- Tier 1: DictConfig / ListConfig — spliced in as-is ---
existing_node = L.call(SimpleEncoder)(depth=4, width=128)
print("Tier 1: existing_node type:", type(existing_node).__name__)

cfg = L.trace(lambda: TieredContainer(node_arg=existing_node, int_arg=42, str_arg="hello"))
print("Result type:", type(cfg).__name__)
print("node_arg type inside cfg:", type(cfg.node_arg).__name__)  # noqa: LACO001  DictConfig
print("int_arg  inside cfg:", cfg.int_arg)    # noqa: LACO001
print("str_arg  inside cfg:", cfg.str_arg)    # noqa: LACO001

# --- Tier 2: OmegaConf primitives — int, float, str, bool, None ---
@L.configurable
@importable
class PrimitiveHolder:
    def __init__(self, i: int, f: float, s: str, b: bool, n):
        pass

cfg_prim = L.trace(lambda: PrimitiveHolder(i=7, f=3.14, s="hi", b=True, n=None))
print("Primitive tier demo:")
print(f"  i={cfg_prim.i}, f={cfg_prim.f}, s={cfg_prim.s!r}, b={cfg_prim.b}, n={cfg_prim.n}")  # noqa: LACO001

# --- Tier 3: arbitrary Python object — raises TypeError ---
@L.configurable
@importable
class WrapperClass:
    def __init__(self, inner):
        pass

live_encoder = Encoder(depth=1)  # real object, NOT a DictConfig

try:
    L.trace(lambda: WrapperClass(inner=live_encoder))
except TypeError as e:
    print("TypeError (Tier 3):")
    # Trim to first 200 chars for readability
    print(str(e)[:200])

@L.configurable
@importable
class DenseLayer:
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        pass

# Pass positional args — no keyword names needed:
cfg_pos = L.trace(lambda: DenseLayer(128, 64))

print("Positional args bound to names:")
print(f"  in_features  = {cfg_pos.in_features}")   # noqa: LACO001  128
print(f"  out_features = {cfg_pos.out_features}")  # noqa: LACO001  64
print(f"  bias         = {cfg_pos.bias}")           # noqa: LACO001  True (filled by apply_defaults)

# apply_defaults() means omitted defaults APPEAR in the node.
# This is important for reproducibility: the saved YAML is fully self-contained.

@L.configurable
@importable
class ConvLayer:
    def __init__(self, in_ch: int, out_ch: int, kernel: int = 3, stride: int = 1, padding: int = 1):
        pass

# Only specify the required args:
cfg_conv = L.trace(lambda: ConvLayer(16, 32))

print("DictConfig with all defaults filled in:")
print(laco.dump(cfg_conv))

@L.configurable
@importable
class InnerEncoder:
    def __init__(self, depth: int, width: int = 256): pass

@L.configurable
@importable
class OuterModel:
    def __init__(self, encoder: InnerEncoder, dim: int = 256): pass

cfg_nested = L.trace(lambda: OuterModel(
    encoder=InnerEncoder(depth=8, width=512),
    dim=1024
))

print("Root type  :", type(cfg_nested).__name__)         # DictConfig
print("encoder type:", type(cfg_nested.encoder).__name__) # noqa: LACO001  DictConfig
print("encoder.depth:", cfg_nested.encoder.depth)         # noqa: LACO001  8
print("encoder.width:", cfg_nested.encoder.width)         # noqa: LACO001  512
print("dim          :", cfg_nested.dim)                    # noqa: LACO001  1024
print()
print("Full YAML:")
print(laco.dump(cfg_nested))

# The nested config can be instantiated. These classes have real bodies (they
# store their attributes) and are @importable so instantiate can resolve them:
@importable
class InnerEncoderReal:
    def __init__(self, depth, width=256):
        self.depth = depth
        self.width = width

@importable
class OuterModelReal:
    def __init__(self, encoder, dim=256):
        self.encoder = encoder
        self.dim = dim

# Re-create the trace pointing to the real classes:
CfgEncoder = L.configurable(InnerEncoderReal)
CfgModel   = L.configurable(OuterModelReal)
cfg_rt = L.trace(lambda: CfgModel(encoder=CfgEncoder(depth=4, width=64), dim=128))

obj = laco.instantiate(cfg_rt)
print(f"Instantiated: {type(obj).__name__}, dim={obj.dim}")              # noqa: LACO001
print(f"  encoder   : {type(obj.encoder).__name__}, depth={obj.encoder.depth}")  # noqa: LACO001

# UndecoratedModel is NOT @L.configurable.
# We use L.call(...) explicitly inside the thunk. It still needs an importable
# target path, so we mark it @importable (but deliberately not @L.configurable):
@importable
class UndecoratedModel:
    def __init__(self, encoder, scale: float = 1.0): pass

@L.configurable
@importable
class DecoEncoder:
    def __init__(self, depth: int): pass

cfg_mixed = L.trace(
    lambda: L.call(UndecoratedModel)(
        encoder=DecoEncoder(depth=4),   # @L.configurable → DictConfig (Tier 1)
        scale=0.5                        # float → Tier 2 primitive
    )
)

print("Mixed trace result:")
print(laco.dump(cfg_mixed))

# What happens if the root is undecorated and not wrapped in L.call?
class RootUndecorated:
    pass

try:
    L.trace(lambda: RootUndecorated())
except TypeError as e:
    print("TypeError (undecorated root):")
    print(str(e)[:200])

# With strict=True the error message is even clearer:
try:
    L.trace(lambda: RootUndecorated(), strict=True)
except TypeError as e:
    print("TypeError (strict=True):")
    print(str(e)[:200])

import pathlib
import tempfile

@L.configurable
@importable
class SimpleNet:
    def __init__(self, layers: int = 3, hidden: int = 256):
        self.layers = layers
        self.hidden = hidden

# Step 1: capture a config via tracing
cfg_original = L.trace(lambda: SimpleNet(layers=5, hidden=128))

print("=== Captured config ===")
yaml_str = laco.dump(cfg_original)
print(yaml_str)

# Step 2: save to a temp file
with tempfile.NamedTemporaryFile(suffix=".yaml", mode="w", delete=False) as fh:
    fh.write(yaml_str)
    tmp_path = pathlib.Path(fh.name)

print(f"Saved to: {tmp_path}")

# Step 3: reload
cfg_reloaded = laco.load(tmp_path)

# Step 4: verify equality
assert laco.dump(cfg_original) == laco.dump(cfg_reloaded)
print("Round-trip YAML is identical.")

# Step 5: instantiate from the reloaded config
obj = laco.instantiate(cfg_reloaded)
print(f"Type    : {type(obj).__name__}")  # SimpleNet
print(f"layers  : {obj.layers}")           # noqa: LACO001  5
print(f"hidden  : {obj.hidden}")           # noqa: LACO001  128

# Cleanup
tmp_path.unlink()