Merge pull request #296 from latchbio/ayush/sdk-ls

latch/types/directory.py

@@ -1,6 +1,8 @@
 from os import PathLike
-from typing import Optional, Type, Union, get_args, get_origin
+from pathlib import Path
+from typing import List, Optional, Type, TypedDict, Union, get_args, get_origin
 
+import gql
 from flytekit.core.annotation import FlyteAnnotation
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
 from flytekit.core.type_engine import TypeEngine, TypeTransformer
@@ -9,13 +11,33 @@
     FlyteDirectory,
     FlyteDirToMultipartBlobTransformer,
 )
+from latch_sdk_gql.execute import execute
+from typing_extensions import Annotated
 
+from latch.types.file import LatchFile
 from latch.types.utils import _is_valid_url
+from latch_cli.utils import urljoins
 
-try:
-    from typing import Annotated
-except ImportError:
-    from typing_extensions import Annotated
+
+class Child(TypedDict):
+    type: str
+    name: str
+
+
+class ChildLdataTreeEdge(TypedDict):
+    child: Child
+
+
+class ChildLdataTreeEdges(TypedDict):
+    nodes: List[ChildLdataTreeEdge]
+
+
+class LDataResolvePathFinalLinkTarget(TypedDict):
+    childLdataTreeEdges: ChildLdataTreeEdges
+
+
+class LdataResolvePathData(TypedDict):
+    finalLinkTarget: LDataResolvePathFinalLinkTarget
 
 
 class LatchDir(FlyteDirectory):
@@ -90,6 +112,51 @@ def downloader():
 
             super().__init__(self.path, downloader, self._remote_directory)
 
+    def iterdir(self) -> List[Union[LatchFile, "LatchDir"]]:
+        ret: List[Union[LatchFile, "LatchDir"]] = []
+
+        if self.remote_path is None:
+            for child in Path(self.path).iterdir():
+                if child.is_dir():
+                    ret.append(LatchDir(str(child)))
+                else:
+                    ret.append(LatchFile(str(child)))
+
+            return ret
+
+        res: Optional[LdataResolvePathData] = execute(
+            gql.gql("""
+            query LDataChildren($argPath: String!) {
+                ldataResolvePathData(argPath: $argPath) {
+                    finalLinkTarget {
+                        childLdataTreeEdges(filter: { child: { removed: { equalTo: false } } }) {
+                            nodes {
+                                child {
+                                    name
+                                    type
+                                }
+                            }
+                        }
+                    }
+                }
+            }"""),
+            {"argPath": self.remote_path},
+        )["ldataResolvePathData"]
+
+        if res is None:
+            raise ValueError(f"No directory found at path: {self}")
+
+        for node in res["finalLinkTarget"]["childLdataTreeEdges"]["nodes"]:
+            child = node["child"]
+
+            path = urljoins(self.remote_path, child["name"])
+            if child["type"] == "DIR":
+                ret.append(LatchDir(path))
+            else:
+                ret.append(LatchFile(path))
+
+        return ret
+
     @property
     def local_path(self) -> str:
         """File path local to the environment executing the task."""

latch/types/file.py

@@ -73,6 +73,7 @@ def __init__(
         if kwargs.get("downloader") is not None:
             super().__init__(self.path, kwargs["downloader"], self._remote_path)
         else:
+
             def downloader():
                 ctx = FlyteContextManager.current_context()
                 if (
@@ -85,15 +86,13 @@ def downloader():
                     local_path_hint = self._remote_path
                     if is_absolute_node_path.match(self._remote_path) is not None:
                         data = execute(
-                            gql.gql(
-                                """
+                            gql.gql("""
                             query getName($argPath: String!) {
                                 ldataResolvePathData(argPath: $argPath) {
                                     name
                                 }
                             }
-                            """
-                            ),
+                            """),
                             {"argPath": self._remote_path},
                         )["ldataResolvePathData"]
 
@@ -159,7 +158,7 @@ def to_python_value(
         ctx: FlyteContext,
         lv: Literal,
         expected_python_type: Union[Type[LatchFile], PathLike],
-    ) -> FlyteFile:
+    ) -> LatchFile:
         uri = lv.scalar.blob.uri
         if expected_python_type is PathLike:
             raise TypeError(

latch_cli/services/cp/autocomplete.py

@@ -17,7 +17,7 @@
     _get_immediate_children_of_node,
     _get_known_domains_for_account,
 )
-from latch_cli.services.cp.path_utils import urljoins
+from latch_cli.utils import urljoins
 
 completion_type = re.compile(
     r"""

latch_cli/services/cp/glob.py

@@ -2,7 +2,7 @@
 from typing import List
 
 from latch_cli.services.cp.ldata_utils import _get_immediate_children_of_node
-from latch_cli.services.cp.path_utils import urljoins
+from latch_cli.utils import urljoins
 
 
 def expand_pattern(remote_path: str) -> List[str]:

latch_cli/services/cp/path_utils.py

@@ -9,53 +9,13 @@
 from latch_cli.services.cp.exceptions import PathResolutionError
 from latch_cli.services.cp.utils import get_auth_header
 
-# todo(ayush): need a better way to check if "latch" has been appended to urllib
-if "latch" not in urllib.parse.uses_netloc:
-    urllib.parse.uses_netloc.append("latch")
-    urllib.parse.uses_relative.append("latch")
-
-
 latch_url_regex = re.compile(r"^(latch)?://")
 
 
 def is_remote_path(path: str) -> bool:
     return latch_url_regex.match(path) is not None
 
 
-def urljoins(*args: str, dir: bool = False) -> str:
-    """Construct a URL by appending paths
-
-    Paths are always joined, with extra `/`s added if missing. Does not allow
-    overriding basenames as opposed to normal `urljoin`. Whether the final
-    path ends in a `/` is still significant and will be preserved in the output
-
-    >>> urljoin("latch:///directory/", "another_directory")
-    latch:///directory/another_directory
-    >>> # No slash means "another_directory" is treated as a filename
-    >>> urljoin(urljoin("latch:///directory/", "another_directory"), "file")
-    latch:///directory/file
-    >>> # Unintentionally overrode the filename
-    >>> urljoins("latch:///directory/", "another_directory", "file")
-    latch:///directory/another_directory/file
-    >>> # Joined paths as expected
-
-    Args:
-        args: Paths to join
-        dir: If true, ensure the output ends with a `/`
-    """
-
-    res = args[0]
-    for x in args[1:]:
-        if res[-1] != "/":
-            res = f"{res}/"
-        res = urljoin(res, x)
-
-    if dir and res[-1] != "/":
-        res = f"{res}/"
-
-    return res
-
-
 scheme = re.compile(
     r"""
     ^(

latch_cli/services/cp/upload.py

@@ -18,14 +18,14 @@
 from latch_cli.constants import latch_constants, units
 from latch_cli.services.cp.config import CPConfig, Progress
 from latch_cli.services.cp.ldata_utils import LDataNodeType, get_node_data
-from latch_cli.services.cp.path_utils import normalize_path, urljoins
+from latch_cli.services.cp.path_utils import normalize_path
 from latch_cli.services.cp.progress import ProgressBarManager, ProgressBars
 from latch_cli.services.cp.utils import (
     get_auth_header,
     get_max_workers,
     human_readable_time,
 )
-from latch_cli.utils import with_si_suffix
+from latch_cli.utils import urljoins, with_si_suffix
 
 if TYPE_CHECKING:
     QueueType: TypeAlias = Queue[Optional[Path]]

latch_cli/utils.py

@@ -3,16 +3,57 @@
 import hashlib
 import os
 import subprocess
+import urllib.parse
 from datetime import datetime, timedelta
 from pathlib import Path
 from typing import List
+from urllib.parse import urljoin
 
 import jwt
 from latch_sdk_config.user import user_config
 
 from latch_cli.constants import latch_constants
 from latch_cli.tinyrequests import get
 
+# todo(ayush): need a better way to check if "latch" has been appended to urllib
+if "latch" not in urllib.parse.uses_netloc:
+    urllib.parse.uses_netloc.append("latch")
+    urllib.parse.uses_relative.append("latch")
+
+
+def urljoins(*args: str, dir: bool = False) -> str:
+    """Construct a URL by appending paths
+
+    Paths are always joined, with extra `/`s added if missing. Does not allow
+    overriding basenames as opposed to normal `urljoin`. Whether the final
+    path ends in a `/` is still significant and will be preserved in the output
+
+    >>> urljoin("latch:///directory/", "another_directory")
+    latch:///directory/another_directory
+    >>> # No slash means "another_directory" is treated as a filename
+    >>> urljoin(urljoin("latch:///directory/", "another_directory"), "file")
+    latch:///directory/file
+    >>> # Unintentionally overrode the filename
+    >>> urljoins("latch:///directory/", "another_directory", "file")
+    latch:///directory/another_directory/file
+    >>> # Joined paths as expected
+
+    Args:
+        args: Paths to join
+        dir: If true, ensure the output ends with a `/`
+    """
+
+    res = args[0]
+    for x in args[1:]:
+        if res[-1] != "/":
+            res = f"{res}/"
+        res = urljoin(res, x)
+
+    if dir and res[-1] != "/":
+        res = f"{res}/"
+
+    return res
+
 
 def retrieve_or_login() -> str:
     """Returns a valid JWT to access Latch, prompting a login flow if needed.