feat: add file type validation and mime type extraction

Dockerfile.dev

@@ -15,7 +15,7 @@ WORKDIR /root
 
 ARG VERSION
 
-RUN apt-get update && apt-get install build-essential vim -y
+RUN apt-get update && apt-get install build-essential vim libmagic1 -y
 
 COPY . /flytekit
 

flytekit/types/file/file.py

@@ -6,6 +6,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 
+import magic
 from dataclasses_json import config
 from marshmallow import fields
 from mashumaro.mixins.json import DataClassJSONMixin
@@ -320,6 +321,31 @@
     def get_literal_type(self, t: typing.Union[typing.Type[FlyteFile], os.PathLike]) -> LiteralType:
         return LiteralType(blob=self._blob_type(format=FlyteFilePathTransformer.get_format(t)))
 
+    def get_mime_type_from_python_type(self, extension: str) -> str:
+        extension_to_mime_type = {
+            "html": "text/html",
+            "jpeg": "image/jpeg",
+            "png": "image/png",
+            "hdf5": "application/x-hdf",
+            "joblib": "application/octet-stream",
+            "pdf": "application/pdf",
+            "python_pickle": "application/octet-stream",
+            "ipynb": "application/x-ipynb+json",
+            "svg": "image/svg+xml",
+            "csv": "text/csv",
+            "onnx": "application/octet-stream",
+            "tfrecord": "application/octet-stream",
+            "txt": "text/plain",
+        }
+        return extension_to_mime_type[extension]
+
+    def validate_file_type(self, python_type, source_path):
+        if FlyteFilePathTransformer.get_format(python_type):
+            real_type = magic.from_file(source_path, mime=True)
+            expected_type = self.get_mime_type_from_python_type(FlyteFilePathTransformer.get_format(python_type))
+            if real_type != expected_type:
+                raise ValueError(f"Incorrect file type, expected {expected_type}, got {real_type}")
+
     def to_literal(
         self,
         ctx: FlyteContext,
@@ -344,6 +370,7 @@
 
         if isinstance(python_val, FlyteFile):
             source_path = python_val.path
+            self.validate_file_type(python_type, source_path)
 
             # If the object has a remote source, then we just convert it back. This means that if someone is just
             # going back and forth between a FlyteFile Python value and a Blob Flyte IDL value, we don't do anything.
@@ -369,6 +396,7 @@
         elif isinstance(python_val, pathlib.Path) or isinstance(python_val, str):
             source_path = str(python_val)
             if issubclass(python_type, FlyteFile):
+                self.validate_file_type(python_type, source_path)
                 if ctx.file_access.is_remote(source_path):
                     should_upload = False
                 else:

setup.py

@@ -75,6 +75,7 @@
         "rich",
         "rich_click",
         "jsonpickle",
+        "python-magic",
     ],
     extras_require=extras_require,
     scripts=[

tests/flytekit/unit/core/test_flyte_file.py

@@ -2,7 +2,7 @@
 import pathlib
 import tempfile
 import typing
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from typing_extensions import Annotated
@@ -19,7 +19,7 @@
 from flytekit.core.workflow import workflow
 from flytekit.models.core.types import BlobType
 from flytekit.models.literals import LiteralMap
-from flytekit.types.file.file import FlyteFile
+from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer
 
 
 # Fixture that ensures a dummy local file
@@ -52,6 +52,79 @@ def my_wf() -> FlyteFile[typing.TypeVar("txt")]:
         assert fh.read() == "Hello World\n"
 
 
+def test_matching_file_types_in_workflow():
+    # TXT
+    @task
+    def t1() -> FlyteFile[typing.TypeVar("txt")]:
+        fname = "/tmp/flytekit_test.txt"
+        with open(fname, "w") as fh:
+            fh.write("Hello World\n")
+        return fname
+
+    @workflow
+    def my_wf() -> FlyteFile[typing.TypeVar("txt")]:
+        f = t1()
+        return f
+
+    res = my_wf()
+    print(type(res))
+    with open(res, "r") as fh:
+        assert fh.read() == "Hello World\n"
+
+
+def test_mismatching_file_types():
+    @task
+    def t1() -> FlyteFile[typing.TypeVar("jpeg")]:
+        fname = "/tmp/flytekit_test.txt"
+        with open(fname, "w") as fh:
+            fh.write("Hello World\n")
+        return fname
+
+    @workflow
+    def my_wf() -> FlyteFile[typing.TypeVar("jpeg")]:
+        f = t1()
+        return f
+
+    with pytest.raises(TypeError) as excinfo:
+        my_wf()
+    assert "Incorrect type, expected image/jpeg, got text/plain" in str(excinfo.value)
+
+
+def test_get_mime_type_from_python_type_success():
+    transformer = TypeEngine.get_transformer(FlyteFile)
+    assert transformer.get_mime_type_from_python_type("html") == "text/html"
+    assert transformer.get_mime_type_from_python_type("jpeg") == "image/jpeg"
+    assert transformer.get_mime_type_from_python_type("png") == "image/png"
+    assert transformer.get_mime_type_from_python_type("hdf5") == "application/x-hdf"
+    assert transformer.get_mime_type_from_python_type("joblib") == "application/octet-stream"
+    assert transformer.get_mime_type_from_python_type("pdf") == "application/pdf"
+    assert transformer.get_mime_type_from_python_type("python_pickle") == "application/octet-stream"
+    assert transformer.get_mime_type_from_python_type("ipynb") == "application/x-ipynb+json"
+    assert transformer.get_mime_type_from_python_type("svg") == "image/svg+xml"
+    assert transformer.get_mime_type_from_python_type("csv") == "text/csv"
+    assert transformer.get_mime_type_from_python_type("onnx") == "application/octet-stream"
+    assert transformer.get_mime_type_from_python_type("tfrecord") == "application/octet-stream"
+    assert transformer.get_mime_type_from_python_type("txt") == "text/plain"
+
+
+def test_get_mime_type_from_python_type_failure():
+    transformer = TypeEngine.get_transformer(FlyteFile)
+    with pytest.raises(KeyError):
+        transformer.get_mime_type_from_python_type("unknown_extension")
+
+
+def test_validate_file_type_incorrect():
+    transformer = TypeEngine.get_transformer(FlyteFile)
+    source_path = "/tmp/flytekit_test.png"
+    source_file_mime_type = "image/png"
+    user_defined_format = "jpeg"
+
+    FlyteFilePathTransformer.get_format = MagicMock(return_value=user_defined_format)
+    with patch("magic.from_file", return_value=source_file_mime_type):
+        with pytest.raises(ValueError, match=f"Incorrect file type, expected image/jpeg, got {source_file_mime_type}"):
+            transformer.validate_file_type(user_defined_format, source_path)
+
+
 def test_file_handling_remote_default_wf_input():
     SAMPLE_DATA = "https://raw.githubusercontent.com/jbrownlee/Datasets/master/pima-indians-diabetes.data.csv"