add class methods, unit tests for flytefile and flytedirectory (#2852)

flytekit/types/directory/types.py

@@ -206,6 +206,23 @@ def new_remote(cls, stem: typing.Optional[str] = None, alt: typing.Optional[str]
         remote_path = ctx.file_access.generate_new_custom_path(alt=alt, stem=stem)
         return cls(path=remote_path)
 
+    @classmethod
+    def new(cls, dirname: str | os.PathLike) -> FlyteFile:
+        """
+        Create a new FlyteDirectory object in current Flyte working directory.
+        """
+
+        if os.path.isabs(dirname):
+            raise ValueError("Path should be relative.")
+
+        ctx = FlyteContextManager.current_context()
+
+        path = os.path.join(ctx.user_space_params.working_directory, dirname)
+
+        os.makedirs(path, exist_ok=False)
+
+        return cls(path=path)
+
     def __class_getitem__(cls, item: typing.Union[typing.Type, str]) -> typing.Type[FlyteDirectory]:
         if item is None:
             return cls
@@ -406,6 +423,13 @@ def __repr__(self):
     def __str__(self):
         return str(self.path)
 
+    def __truediv__(self, other: str | os.PathLike) -> Path:
+        """
+        This is a convenience method to allow for easy concatenation of paths.
+        """
+
+        return Path(self.path) / other
+
 
 class FlyteDirToMultipartBlobTransformer(AsyncTypeTransformer[FlyteDirectory]):
     """

flytekit/types/file/file.py

@@ -214,6 +214,21 @@ def from_source(cls, source: str | os.PathLike) -> FlyteFile:
         t = FlyteFilePathTransformer()
         return t.to_python_value(ctx, lit, cls)
 
+    @classmethod
+    def new(cls, filename: str | os.PathLike) -> FlyteFile:
+        """
+        Create a new FlyteFile object in the current Flyte working directory
+        """
+
+        if os.path.isabs(filename):
+            raise ValueError("Path should be relative.")
+
+        ctx = FlyteContextManager.current_context()
+
+        path = os.path.join(ctx.user_space_params.working_directory, filename)
+
+        return cls(path=path)
+
     def __class_getitem__(cls, item: typing.Union[str, typing.Type]) -> typing.Type[FlyteFile]:
         from flytekit.types.file import FileExt
 

tests/flytekit/unit/types/directory/test_types.py

@@ -1,6 +1,6 @@
 import mock
 
-from flytekit import FlyteContext
+from flytekit import FlyteContext, FlyteContextManager
 from flytekit.types.directory import FlyteDirectory
 from flytekit.types.file import FlyteFile
 
@@ -27,6 +27,16 @@ def test_new_remote_dir_alt():
     assert "my-alt-bucket" in ff.path
     assert "my-stem" in ff.path
 
+def test_new_auto_new_dir():
+    fd = FlyteDirectory.new("my_dir")
+    assert FlyteContextManager.current_context().user_space_params.working_directory in fd.path
+
+def test_add_path_to_dir():
+    fd = FlyteDirectory.new("my_other_dir")
+    cwd = FlyteContextManager.current_context().user_space_params.working_directory
+    assert cwd in str(fd / "myfile.txt")
+
+
 @mock.patch("flytekit.types.directory.types.os.name", "nt")
 def test_sep_nt():
     fd = FlyteDirectory(path="file://mypath")

tests/flytekit/unit/types/file/test_types.py

@@ -5,3 +5,8 @@ def test_new_remote_alt():
     ff = FlyteFile.new_remote_file(alt="my-alt-prefix", name="my-file.txt")
     assert "my-alt-prefix" in ff.path
     assert "my-file.txt" in ff.path
+
+def test_new_auto_file():
+    ff = FlyteFile.new("my-file.txt")
+    cwd = FlyteContextManager.current_context().user_space_params.working_directory
+    assert cwd in ff.path