diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index cc7ba66bed..44d2f112d8 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -179,13 +179,21 @@ def extension(cls) -> str: return "" @classmethod - def new_remote_file(cls, name: typing.Optional[str] = None) -> FlyteFile: + def new_remote_file(cls, alt: typing.Optional[str] = None, name: typing.Optional[str] = None) -> FlyteFile: """ Create a new FlyteFile object with a remote path. + + :param alt: If you want to specify a different prefix head than the default one, you can specify it here. + :param name: If you want to specify a different name for the file, you can specify it here. """ ctx = FlyteContextManager.current_context() r = name or ctx.file_access.get_random_string() - remote_path = ctx.file_access.join(ctx.file_access.raw_output_prefix, r) + pref = ctx.file_access.raw_output_prefix + if alt: + s_pref = pref.split("/") + s_pref[2] = alt + pref = "/".join(s_pref) + remote_path = ctx.file_access.join(pref, r) return cls(path=remote_path) @classmethod diff --git a/tests/flytekit/unit/types/file/test_types.py b/tests/flytekit/unit/types/file/test_types.py new file mode 100644 index 0000000000..6d37176833 --- /dev/null +++ b/tests/flytekit/unit/types/file/test_types.py @@ -0,0 +1,6 @@ +from flytekit.types.file import FlyteFile +from flytekit import FlyteContextManager + +def test_new_remote_alt(): + ff = FlyteFile.new_remote_file(alt="my-alt-prefix") + assert "my-alt-prefix" in ff.path