add optional support for pydantic (fsspec#395)

* add optional support for pydantic

* fix type checking

* try abspath for windows

* more mypy ignores

* Revert "try abspath for windows"

This reverts commit c0090c5.

* try fix windows

* refactor imports

* parametrize tests

Author: pmeier

pyproject.toml

@@ -41,6 +41,7 @@ tests = [
     "pytest-mock >=3.12.0",
     "pylint >=2.17.4",
     "mypy >=1.10.0",
+    "pydantic >=2",
     "pytest-mypy-plugins >=3.1.2",
     "packaging",
 ]
@@ -168,6 +169,22 @@ ignore_missing_imports = true
 module = "smbprotocol.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "pydantic.*"
+ignore_errors = true
+
+[[tool.mypy.overrides]]
+module = "pydantic_core.*"
+ignore_errors = true
+
+[[tool.mypy.overrides]]
+module = "typing_inspection.*"
+ignore_errors = true
+
+[[tool.mypy.overrides]]
+module = "annotated_types.*"
+ignore_errors = true
+
 [tool.pylint.format]
 max-line-length = 88
 

upath/core.py

@@ -44,6 +44,9 @@
     else:
         from typing_extensions import Self
 
+    from pydantic import GetCoreSchemaHandler
+    from pydantic_core.core_schema import CoreSchema
+
 
 __all__ = ["UPath"]
 
@@ -934,3 +937,64 @@ def match(self, pattern: str) -> bool:
         if not pattern:
             raise ValueError("pattern cannot be empty")
         return self.full_match(pattern.replace("**", "*"))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, _source_type: Any, _handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        from pydantic_core import core_schema
+
+        deserialization_schema = core_schema.chain_schema(
+            [
+                core_schema.no_info_plain_validator_function(
+                    lambda v: {"path": v} if isinstance(v, str) else v,
+                ),
+                core_schema.typed_dict_schema(
+                    {
+                        "path": core_schema.typed_dict_field(
+                            core_schema.str_schema(), required=True
+                        ),
+                        "protocol": core_schema.typed_dict_field(
+                            core_schema.with_default_schema(
+                                core_schema.str_schema(), default=""
+                            ),
+                            required=False,
+                        ),
+                        "storage_options": core_schema.typed_dict_field(
+                            core_schema.with_default_schema(
+                                core_schema.dict_schema(
+                                    core_schema.str_schema(),
+                                    core_schema.any_schema(),
+                                ),
+                                default_factory=dict,
+                            ),
+                            required=False,
+                        ),
+                    },
+                    extra_behavior="forbid",
+                ),
+                core_schema.no_info_plain_validator_function(
+                    lambda dct: cls(
+                        dct.pop("path"),
+                        protocol=dct.pop("protocol"),
+                        **dct["storage_options"],
+                    )
+                ),
+            ]
+        )
+
+        serialization_schema = core_schema.plain_serializer_function_ser_schema(
+            lambda u: {
+                "path": u.path,
+                "protocol": u.protocol,
+                "storage_options": dict(u.storage_options),
+            }
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=deserialization_schema,
+            python_schema=core_schema.union_schema(
+                [core_schema.is_instance_schema(UPath), deserialization_schema]
+            ),
+            serialization=serialization_schema,
+        )

upath/tests/test_pydantic.py

@@ -0,0 +1,113 @@
+import json
+from os.path import abspath
+
+import pydantic
+import pydantic_core
+import pytest
+from fsspec.implementations.http import get_client
+
+from upath import UPath
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/abc",
+        "file:///abc",
+        "memory://abc",
+        "s3://bucket/key",
+        "https://www.example.com",
+    ],
+)
+@pytest.mark.parametrize("source", ["json", "python"])
+def test_validate_from_str(path, source):
+    expected = UPath(path)
+
+    ta = pydantic.TypeAdapter(UPath)
+    if source == "json":
+        actual = ta.validate_json(json.dumps(path))
+    else:  # source == "python"
+        actual = ta.validate_python(path)
+
+    assert abspath(actual.path) == abspath(expected.path)
+    assert actual.protocol == expected.protocol
+
+
+@pytest.mark.parametrize(
+    "dct",
+    [
+        {
+            "path": "/my/path",
+            "protocol": "file",
+            "storage_options": {"foo": "bar", "baz": 3},
+        }
+    ],
+)
+@pytest.mark.parametrize("source", ["json", "python"])
+def test_validate_from_dict(dct, source):
+    ta = pydantic.TypeAdapter(UPath)
+    if source == "json":
+        output = ta.validate_json(json.dumps(dct))
+    else:  # source == "python"
+        output = ta.validate_python(dct)
+
+    assert abspath(output.path) == abspath(dct["path"])
+    assert output.protocol == dct["protocol"]
+    assert output.storage_options == dct["storage_options"]
+
+
+@pytest.mark.parametrize(
+    "path",
+    [
+        "/abc",
+        "file:///abc",
+        "memory://abc",
+        "s3://bucket/key",
+        "https://www.example.com",
+    ],
+)
+def test_validate_from_instance(path):
+    input = UPath(path)
+
+    output = pydantic.TypeAdapter(UPath).validate_python(input)
+
+    assert output is input
+
+
+@pytest.mark.parametrize(
+    ("args", "kwargs"),
+    [
+        (
+            ("/my/path",),
+            {
+                "protocol": "file",
+                "foo": "bar",
+                "baz": 3,
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize("mode", ["json", "python"])
+def test_dump(args, kwargs, mode):
+    u = UPath(*args, **kwargs)
+
+    output = pydantic.TypeAdapter(UPath).dump_python(u, mode=mode)
+
+    assert output["path"] == u.path
+    assert output["protocol"] == u.protocol
+    assert output["storage_options"] == u.storage_options
+
+
+def test_dump_non_serializable_python():
+    output = pydantic.TypeAdapter(UPath).dump_python(
+        UPath("https://www.example.com", get_client=get_client), mode="python"
+    )
+
+    assert output["storage_options"]["get_client"] is get_client
+
+
+def test_dump_non_serializable_json():
+    with pytest.raises(pydantic_core.PydanticSerializationError, match="unknown type"):
+        pydantic.TypeAdapter(UPath).dump_python(
+            UPath("https://www.example.com", get_client=get_client), mode="json"
+        )