Merge pull request #390 from latchbio/rahuldesai1/optional-params

fix: optional parameters for snakemake

latch_cli/snakemake/config/utils.py

@@ -220,11 +220,13 @@ def type_repr(t: Type, *, add_namespace: bool = False) -> str:
     if get_origin(t) is Union:
         args = get_args(t)
 
-        assert len(args) > 0
+        if len(args) != 2 or args[1] is not type(None):
+            raise ValueError(
+                "Union types other than Optional are not yet supported in Snakemake"
+                " workflows."
+            )
 
-        return (
-            f"typing.Union[{', '.join([type_repr(arg, add_namespace=add_namespace) for arg in args])}]"
-        )
+        return f"typing.Optional[{type_repr(args[0], add_namespace=add_namespace)}]"
 
     if get_origin(t) is Annotated:
         args = get_args(t)

latch_cli/snakemake/workflow.py

@@ -397,7 +397,6 @@ def get_file_download_code(
         param_path: str,
     ) -> str:
         code_block = ""
-        code_block += f"if {param} is not None:\n"
 
         param_name = identifier_from_str(param)
         touch_str = f"{param}._create_imposters()"
@@ -407,42 +406,30 @@ def get_file_download_code(
                 f" Path({param}).resolve()"
             )
 
-        code_block += reindent(
-            rf"""
+        code_block += dedent(rf"""
             {param_name}_dst_p = Path("{path}")
-
             {touch_str}
             {param_name}_p = Path({param}.path)
             print(f"  {{file_name_and_size({param_name}_p)}}")
-            """,
-            1,
-        )
+            """)
 
         if t is LatchDir:
-            code_block += f"    {param_name}_p.iterdir()\n"
+            code_block += f"{param_name}_p.iterdir()\n"
 
-        code_block += reindent(
-            rf"""
+        code_block += dedent(rf"""
             print(f"Moving {param} to {{{param_name}_dst_p}}")
-
             update_mapping({param_name}_p, {param_name}_dst_p, {param}.remote_path, local_to_remote_path_mapping)
             check_exists_and_rename(
                 {param_name}_p,
                 {param_name}_dst_p
             )
-
-            """,
-            1,
-        )
+            """)
 
         if config:
-            code_block += reindent(
-                rf"""
+            code_block += dedent(rf"""
                 print(f"Saving parameter value {param} = {str(path)}")
                 overwrite_config[{self._param_path_str(param_path)}] = {repr(str(path))}
-                """,
-                1,
-            )
+                """)
 
         return code_block
 
@@ -453,6 +440,14 @@ def get_param_code(
         file_meta: FileMetadata,
         param_path: List[Union[str, int]],
     ) -> str:
+        if get_origin(t) is Union:
+            args = get_args(t)
+            assert len(args) > 0
+            code_block = f"if {param} is not None:\n"
+            return code_block + reindent(
+                self.get_param_code(args[0], param, file_meta, param_path), 1
+            )
+
         param_meta = self.parameter_metadata.get(param)
         # support workflows that use the legacy SnakemakeFileParameter
         if isinstance(param_meta, SnakemakeFileParameter):
@@ -474,20 +469,17 @@ def get_param_code(
             return code_block
 
         if file_meta is None or file_meta.get(param) is None or is_primitive_type(t):
-            return dedent(rf"""
-                print(f"Saving parameter value {param} = get_parameter_json_value({param})")
-                overwrite_config[{self._param_path_str(param_path)}] = get_parameter_json_value({param})
+            return dedent(f"""
+                if {param} is not None:
+                    print(f"Saving parameter value {param} = get_parameter_json_value({param})")
+                    overwrite_config[{self._param_path_str(param_path)}] = get_parameter_json_value({param})
                 """)
 
         if get_origin(t) is Annotated:
             args = get_args(t)
             assert len(args) > 0
             return self.get_param_code(args[0], param, file_meta, param_path)
 
-        if get_origin(t) is Union:
-            args = get_args(t)
-            return self.get_param_code(args[0], param, file_meta, param_path)
-
         meta = file_meta.get(param)
         if t in (LatchFile, LatchDir):
             return self.get_file_download_code(
@@ -567,7 +559,7 @@ def get_fn_code(
 
         for param, t in self.python_interface.inputs.items():
             code_block += reindent(
-                self.get_param_code(t, param, self.file_metadata, [param]),
+                self.get_param_code(t, param, self.file_metadata, [param]) + "\n",
                 1,
             )