work in progress

Signed-off-by: Ketan Umare <ketan.umare@gmail.com>

flytekit/types/directory/types.py

@@ -3,6 +3,7 @@
 import os
 import pathlib
 import typing
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -182,6 +183,19 @@ def remote_source(self) -> str:
     def download(self) -> str:
         return self.__fspath__()
 
+    @contextmanager
+    def open(self, mode: str, **kwargs):
+        try:
+            import fsspec
+            final_path = self.remote_path if self.remote_path else self.path
+            open_files: fsspec.core.OpenFiles = fsspec.open_files(final_path, mode, **kwargs)
+            # TODO should we wrap this up into FlyteFile?
+            return open_files
+        except ImportError as e:
+            print("To use streaming files, please install fsspec."
+                  " Note: This will be bundled with flytekit in the future.")
+            raise
+
     def __repr__(self):
         return self.path
 

flytekit/types/file/file.py

@@ -239,7 +239,7 @@ def open(self, mode: str, **kwargs):
         try:
             import fsspec
             final_path = self.remote_path if self.remote_path else self.path
-            open_file: fsspec.core.OpenFile = fsspec.open(final_path, mode)
+            open_file: fsspec.core.OpenFile = fsspec.open(final_path, mode, **kwargs)
             try:
                 yield open_file.open()
             finally:

tests/flytekit/unit/clients/test_raw.py

@@ -241,4 +241,4 @@ def test__handle_invalid_create_request_decorator_raises(mock_to_JSON, mock_logg
     client._stub.CreateWorkflow = mock.Mock(side_effect=err)
     with pytest.raises(grpc.RpcError):
         _handle_invalid_create_request(client.create_workflow("/flyteidl.service.AdminService/CreateWorkflow"))
-    mock_logger.error.assert_called_with("There is already a workflow with different structure.")
+    mock_logger.error.assert_called_with("There is already a workflow with different structure.")