From ba9b4b886642470a9fe82bc4eb32ddf7c3003a62 Mon Sep 17 00:00:00 2001
From: granthamtaylor <granthamtaylor@icloud.com>
Date: Wed, 23 Oct 2024 09:31:39 -0400
Subject: [PATCH] add class methods, unit tests for flytefile and
 flytedirectory

---
 flytekit/types/directory/types.py             | 24 +++++++++++++++++++
 flytekit/types/file/file.py                   | 15 ++++++++++++
 .../unit/types/directory/test_types.py        | 12 +++++++++-
 tests/flytekit/unit/types/file/test_types.py  |  5 ++++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py
index 52249e2977..7e22879126 100644
--- a/flytekit/types/directory/types.py
+++ b/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]):
     """
diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py
index bf08a1f535..a76e5f4e05 100644
--- a/flytekit/types/file/file.py
+++ b/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
 
diff --git a/tests/flytekit/unit/types/directory/test_types.py b/tests/flytekit/unit/types/directory/test_types.py
index 1b9cf4be97..36f214a9da 100644
--- a/tests/flytekit/unit/types/directory/test_types.py
+++ b/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")
diff --git a/tests/flytekit/unit/types/file/test_types.py b/tests/flytekit/unit/types/file/test_types.py
index 7cc6e42fea..cc3e575428 100644
--- a/tests/flytekit/unit/types/file/test_types.py
+++ b/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