Loading, Saving, and the CLI

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

print(f"laco version: {laco.__version__ if hasattr(laco, '__version__') else 'installed'}")

configs://examples/mlp.py ? hps.dim_hidden=512 & hps.num_layers=4 # model
└─── file path ──────────┘ └────── overrides (query string) ──────┘ └ key ┘

# Form 1: URL with query string and fragment — everything in one string.
# The query string becomes override strings; the fragment becomes the key.
cfg1 = laco.load(
    "configs://examples/linear_regression.py?hps.in_features=16#model"
)
print("Form 1 — cfg1 (model node, in_features=16):")
print(OmegaConf.to_yaml(cfg1))

# Form 2: path + positional override strings + keyword key.
# Cleaner when you need to compute overrides dynamically.
cfg2 = laco.load(
    "configs://examples/linear_regression.py",
    "hps.in_features=16",
    key="model",
)
print("Form 2 — cfg2 (identical to cfg1):")
print(OmegaConf.to_yaml(cfg2))

# Sanity-check: both forms produce the same config
assert OmegaConf.to_yaml(cfg1) == OmegaConf.to_yaml(cfg2), "Mismatch!"
print("Forms 1 and 2 are identical.")

# Form 3: bare path — loads the entire module namespace.
# The result is a DictConfig with .model, .optimizer, .hps keys
# (whatever is listed in __all__ of the config file).
cfg3 = laco.load("configs://examples/linear_regression.py")
print("Form 3 — top-level keys:", list(cfg3.keys()))
print()
print(OmegaConf.to_yaml(cfg3))

import pathlib
import tempfile

# 1. Load the model sub-config from the linear_regression example.
cfg = laco.load("configs://examples/linear_regression.py#model")
print("Loaded config type:", type(cfg).__name__)
print()

# 2. Serialise to a YAML string.
yaml_str = laco.dump(cfg)
print("YAML output:")
print(yaml_str)

# 3. Write to a temp file and reload — the round-trip.
with tempfile.NamedTemporaryFile(suffix=".yaml", mode="w", delete=False) as f:
    f.write(yaml_str)
    tmp_path = f.name

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

# Reload from the .yaml file — laco detects the extension automatically.
cfg_reloaded = laco.load(tmp_path)
print("Reloaded config type:", type(cfg_reloaded).__name__)
print()
print("Round-tripped YAML:")
print(OmegaConf.to_yaml(cfg_reloaded))

# Cleanup
pathlib.Path(tmp_path).unlink()
print("Temp file removed.")

# laco.save is a thin convenience wrapper:
# it calls dump(), writes the file, then reloads it (unless mode=SaveMode.NO_RELOAD).
# The returned object is the freshly loaded DictConfig.

from laco._io import SaveMode

with tempfile.NamedTemporaryFile(suffix=".yaml", delete=False) as f:
    save_path = f.name

cfg_from_save = laco.save(cfg, save_path)
print("laco.save() returned type:", type(cfg_from_save).__name__)
print("Keys:", list(cfg_from_save.keys()))

pathlib.Path(save_path).unlink()
print("Temp file removed.")

laco.load(path)  -->  DictConfig  --laco.dump()-->  YAML string
                                                         |
                          <-------- laco.load() ---------+

[project.entry-points."expath.protocols"]
configs = "laco.examples:__path__[0]"

laco.load("path/to/my_config.py")
laco.load("/absolute/path/to/config.py")
laco.load("/absolute/path/to/saved.yaml")

# Resolve the actual path on disk using expath.
try:
    import expath
    resolved = expath.resolve("configs://examples/mlp.py")
    print("configs://examples/mlp.py  resolves to:")
    print(" ", resolved)
except Exception as e:
    print(f"expath not available or resolution failed: {e}")

# Demonstrate loading from an absolute filesystem path.
try:
    import expath
    abs_path = expath.resolve("configs://examples/linear_regression.py")
    cfg_from_abs = laco.load(str(abs_path))
    print("Loaded from absolute path — top-level keys:", list(cfg_from_abs.keys()))
except Exception as e:
    print(f"Could not demonstrate absolute path loading: {e}")

# Compose: load and dump the entire linear_regression config as YAML.
# The output is identical to laco.dump(laco.load(...)).
!laco compose configs://examples/linear_regression.py

# Compose with an override — same as laco.load(..., "hps.in_features=16").
!laco compose configs://examples/linear_regression.py hps.in_features=16

# Show: pretty-print the config (equivalent to compose with no overrides).
!laco show configs://examples/mlp.py

# Alternative: subprocess.run — useful when you want to capture the output
# in Python code without relying on shell magic.
import subprocess

result = subprocess.run(
    ["laco", "compose", "configs://examples/linear_regression.py"],
    capture_output=True,
    text=True,
)

if result.returncode == 0:
    print("Exit code: 0 (success)")
    print("First 500 characters of output:")
    print(result.stdout[:500])
else:
    print("Exit code:", result.returncode)
    print("stderr:", result.stderr)

# python -m laco.cli is available when the laco CLI is not on PATH.
result2 = subprocess.run(
    ["python", "-m", "laco.cli", "compose", "configs://examples/linear_regression.py"],
    capture_output=True,
    text=True,
)
print("python -m laco.cli exit code:", result2.returncode)
assert result2.stdout == result.stdout, "Output differs!"
print("Output is identical to the laco CLI.")

from laco._io import LoadMode, DEFAULT_SAFE_ALLOWED_MODULES

print("DEFAULT_SAFE_ALLOWED_MODULES:", DEFAULT_SAFE_ALLOWED_MODULES)

# Load in SAFE mode — the linear_regression example only imports laco and torch,
# so it is fully compatible with the default allowlist.
cfg_safe = laco.load(
    "configs://examples/linear_regression.py",
    mode=LoadMode.SAFE,
)
print("SAFE mode loaded successfully. Keys:", list(cfg_safe.keys()))

# You can widen (or narrow) the allowlist with allowed_modules=.
# Example: also allow 'numpy' in addition to the defaults.
cfg_safe_wide = laco.load(
    "configs://examples/linear_regression.py",
    mode=LoadMode.SAFE,
    allowed_modules=["laco", "omegaconf", "torch", "torch.nn", "numpy"],
)
print("Wide-allowlist SAFE mode loaded successfully.")

# Demonstrate that SAFE mode rejects disallowed imports.
import textwrap
import pathlib
import tempfile

# Write a config that tries to import 'os' — not in the allowlist.
evil_config = textwrap.dedent("""
    import os
    import laco.language as L
    from torch import nn

    model = L.call(nn.Linear)(in_features=8, out_features=1)
""")

with tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False) as f:
    f.write(evil_config)
    evil_path = f.name

try:
    laco.load(evil_path, mode=LoadMode.SAFE)
    print("ERROR: should have raised a LacoLoadError!")
except laco.LacoLoadError as e:
    print("SAFE mode correctly blocked the import:")
    print(" ", str(e)[:120])
finally:
    pathlib.Path(evil_path).unlink()