feat: add support for nested file paths in latch metadata

latch/types/metadata.py

@@ -380,6 +380,8 @@ class _IsDataclass(Protocol):
     Type[Enum],
     Type[_IsDataclass],
     Type[List["ParameterType"]],
+    Type[LatchFile],
+    Type[LatchDir],
 ]
 
 
@@ -393,6 +395,7 @@ class SnakemakeParameter(LatchParameter):
     default: Optional[Any] = None
 
 
+# DEPRECATED: use `file_metadata` keyword in `SnakemakeMetadata` instead
 @dataclass
 class SnakemakeFileParameter(SnakemakeParameter):
     type: Optional[
@@ -418,6 +421,22 @@ class SnakemakeFileParameter(SnakemakeParameter):
     """
 
 
+@dataclass
+class SnakemakeFileMetadata:
+    path: Optional[Path] = None
+    """
+    The path where the file passed to this parameter will be copied.
+    """
+    config: bool = False
+    """
+    Whether or not the file path is exposed in the Snakemake config
+    """
+    download: bool = False
+    """
+    Whether or not the file is downloaded in the JIT step
+    """
+
+
 @dataclass
 class LatchMetadata:
     """Class for organizing workflow metadata
@@ -549,6 +568,9 @@ class EnvironmentConfig:
     """
 
 
+FileMetadata: TypeAlias = Dict[str, Union[SnakemakeFileMetadata, "FileMetadata"]]
+
+
 @dataclass
 class SnakemakeMetadata(LatchMetadata):
     """Class for organizing Snakemake workflow metadata"""
@@ -573,6 +595,10 @@ class SnakemakeMetadata(LatchMetadata):
     """
     A dictionary mapping parameter names (strings) to `SnakemakeParameter` objects
     """
+    file_metadata: FileMetadata = field(default_factory=dict)
+    """
+    A dictionary mapping parameter names to `SnakemakeFileMetadata` objects
+    """
 
     def __post_init__(self):
         if self.name is None:

latch_cli/snakemake/config/parser.py

@@ -1,3 +1,4 @@
+from dataclasses import fields, is_dataclass
 from pathlib import Path
 from typing import Dict, List, Tuple, Type, TypeVar
 
@@ -6,11 +7,19 @@
 
 from latch.types.directory import LatchDir
 from latch.types.file import LatchFile
-from latch_cli.snakemake.workflow import reindent
+from latch_cli.snakemake.utils import reindent
 from latch_cli.utils import identifier_from_str
 
 from ..serialize_utils import best_effort_display_name
-from .utils import JSONValue, get_preamble, parse_type, parse_value, type_repr
+from .utils import (
+    JSONValue,
+    get_preamble,
+    is_list_type,
+    is_primitive_type,
+    parse_type,
+    parse_value,
+    type_repr,
+)
 
 T = TypeVar("T")
 
@@ -57,13 +66,60 @@ def parse_config(
     parsed: Dict[str, Type] = {}
     for k, v in res.items():
         typ = parse_type(v, k, infer_files=infer_files)
-        val = parse_value(typ, v)
+        val, default = parse_value(typ, v)
 
-        parsed[k] = (typ, val)
+        parsed[k] = (typ, (val, default))
 
     return parsed
 
 
+def file_metadata_str(typ: Type, value: JSONValue, level: int = 0):
+    if is_primitive_type(typ):
+        return None
+
+    if typ in {LatchFile, LatchDir}:
+        return reindent(
+            f"""\
+            SnakemakeFileMetadata(
+                path={repr(value)},
+                config=True,
+            ),\n""",
+            level,
+        )
+
+    metadata: List[str] = []
+    if is_list_type(typ):
+        template = """
+        [
+        __metadata__],\n"""
+        for val in value:
+            metadata_str = file_metadata_str(typ.__args__[0], val, level + 1)
+            if metadata_str is None:
+                continue
+            metadata.append(metadata_str)
+    else:
+        template = """
+        {
+        __metadata__},\n"""
+        assert is_dataclass(typ)
+        for field in fields(typ):
+            metadata_str = file_metadata_str(
+                field.type, value.__getattribute__(field.name), level
+            )
+            if metadata_str is None:
+                continue
+            metadata_str = f"{repr(identifier_from_str(field.name))}: {metadata_str}"
+            metadata.append(reindent(metadata_str, level + 1))
+
+    if len(metadata) == 0:
+        return None
+
+    return reindent(
+        template,
+        level,
+    ).replace("__metadata__", "".join(metadata), level + 1)
+
+
 # todo(ayush): print informative stuff here ala register
 def generate_metadata(
     config_path: Path,
@@ -76,37 +132,35 @@ def generate_metadata(
 
     preambles: List[str] = []
     params: List[str] = []
+    file_metadata: List[str] = []
 
-    for k, (typ, val) in parsed.items():
+    for k, (typ, (val, default)) in parsed.items():
         preambles.append(get_preamble(typ))
 
-        is_file = typ in {LatchFile, LatchDir}
-        param_typ = "SnakemakeFileParameter" if is_file else "SnakemakeParameter"
-
         param_str = reindent(
             f"""\
-            {repr(identifier_from_str(k))}: {param_typ}(
+            {repr(identifier_from_str(k))}: SnakemakeParameter(
                 display_name={repr(best_effort_display_name(k))},
                 type={type_repr(typ)},
-            __config____default__),""",
+            __default__),""",
             0,
         )
 
-        config = ""
-        if is_file:
-            config = "    config=True,\n"
-
-        param_str = param_str.replace("__config__", config)
-
-        default = ""
-        if generate_defaults and val is not None:
-            default = f"    default={repr(val)},\n"
+        default_str = ""
+        if generate_defaults and default is not None:
+            default_str = f"    default={repr(default)},\n"
 
-        param_str = param_str.replace("__default__", default)
+        param_str = param_str.replace("__default__", default_str)
 
         param_str = reindent(param_str, 1)
         params.append(param_str)
 
+        metadata_str = file_metadata_str(typ, val)
+        if metadata_str is None:
+            continue
+        metadata_str = f"{repr(identifier_from_str(k))}: {metadata_str}"
+        file_metadata.append(reindent(metadata_str, 1))
+
     metadata_root = Path("latch_metadata")
     if metadata_root.is_file():
         if not click.confirm("A file exists at `latch_metadata`. Delete it?"):
@@ -146,7 +200,7 @@ def generate_metadata(
                 from latch.types.metadata import SnakemakeMetadata, LatchAuthor
                 from latch.types.directory import LatchDir
 
-                from .parameters import generated_parameters
+                from .parameters import generated_parameters, file_metadata
 
                 SnakemakeMetadata(
                     output_dir=LatchDir("latch:///your_output_directory"),
@@ -156,6 +210,7 @@ def generate_metadata(
                     ),
                     # Add more parameters
                     parameters=generated_parameters,
+                    file_metadata=file_metadata,
                 )
                 """,
                 0,
@@ -178,22 +233,27 @@ def generate_metadata(
             from dataclasses import dataclass
             import typing
 
-            from latch.types.metadata import SnakemakeParameter, SnakemakeFileParameter
+            from latch.types.metadata import SnakemakeParameter, SnakemakeFileParameter, SnakemakeFileMetadata
             from latch.types.file import LatchFile
             from latch.types.directory import LatchDir
 
             __preambles__
 
             # Import these into your `__init__.py` file:
             #
-            # from .parameters import generated_parameters
+            # from .parameters import generated_parameters, file_metadata
 
             generated_parameters = {
             __params__
             }
+
+            file_metadata = {
+            __file_metadata__}
+
             """,
             0,
         )
         .replace("__preambles__", "".join(preambles))
         .replace("__params__", "\n".join(params))
+        .replace("__file_metadata__", "".join(file_metadata))
     )

latch_cli/snakemake/config/utils.py

@@ -85,6 +85,8 @@
     "hdf5",
     "he5",
     "h5ad",
+    "yaml",
+    "yml",
 }
 
 
@@ -108,11 +110,13 @@ def parse_type(
             parse_type(
                 x,
                 name,
-                infer_files=False,  # todo(ayush): enable recursive file inference
+                infer_files=infer_files,
             )
             for x in v
         )
-        return List[Union[parsed_types]]
+        if len(set(parsed_types)) > 1:
+            raise ValueError("all types in List must be same")
+        return List[parsed_types[0]]
 
     assert isinstance(v, dict)
 
@@ -124,44 +128,47 @@ def parse_type(
         fields[identifier_from_str(k)] = parse_type(
             x,
             k,
-            infer_files=False,  # todo(ayush): enable recursive file inference
+            infer_files=infer_files,
         )
 
     return make_dataclass(identifier_from_str(name), fields.items())
 
 
+# returns raw value and generated default
 def parse_value(t: Type, v: JSONValue):
     if v is None:
-        return None
+        return None, None
 
     if t in {LatchFile, LatchDir}:
         # ayush: autogenerated defaults don't make sense for files/dirs since their
         # value in the config is their local path
-        return None
+        return v, None
 
     if is_primitive_value(v):
-        return v
+        return v, v
 
     if isinstance(v, list):
         assert get_origin(t) is list
 
         sub_type = get_args(t)[0]
-
-        return [parse_value(sub_type, x) for x in v]
+        res = [parse_value(sub_type, x) for x in v]
+        return [x[0] for x in res], [x[1] for x in res]
 
     assert isinstance(v, dict), v
     assert is_dataclass(t), t
 
     ret = {}
+    defaults = {}
     fs = {identifier_from_str(f.name): f for f in fields(t)}
 
     for k, x in v.items():
         sanitized = identifier_from_str(k)
         assert sanitized in fs, sanitized
+        val, default = parse_value(fs[sanitized].type, x)
+        ret[sanitized] = val
+        defaults[sanitized] = default
 
-        ret[sanitized] = parse_value(fs[sanitized].type, x)
-
-    return t(**ret)
+    return t(**ret), t(**defaults)
 
 
 def is_primitive_type(
@@ -174,6 +181,10 @@ def is_primitive_value(val: object) -> TypeGuard[Union[None, str, bool, int, flo
     return is_primitive_type(type(val))
 
 
+def is_list_type(typ: Type) -> TypeGuard[Type[List]]:
+    return get_origin(typ) is list
+
+
 def type_repr(t: Type, *, add_namespace: bool = False) -> str:
     if is_primitive_type(t) or t is LatchFile or t is LatchDir:
         return t.__name__

latch_cli/snakemake/utils.py

@@ -1,3 +1,4 @@
+import textwrap
 from pathlib import Path
 from typing import Optional
 
@@ -16,3 +17,10 @@ def load_snakemake_metadata(pkg_root: Path) -> Optional[Path]:
         import_module_by_path(old_meta)
 
         return old_meta
+
+
+# todo(maximsmol): use a stateful writer that keeps track of indent level
+def reindent(x: str, level: int) -> str:
+    if x[0] == "\n":
+        x = x[1:]
+    return textwrap.indent(textwrap.dedent(x), "    " * level)