support JSON schema generation (#451)

Author: pmeier

upath/core.py

@@ -1387,38 +1387,50 @@ def __get_pydantic_core_schema__(
     ) -> CoreSchema:
         from pydantic_core import core_schema
 
-        deserialization_schema = core_schema.chain_schema(
+        str_schema = core_schema.chain_schema(
             [
+                core_schema.str_schema(),
                 core_schema.no_info_plain_validator_function(
-                    lambda v: {"path": v} if isinstance(v, str) else v,
+                    lambda path: {
+                        "path": path,
+                        "protocol": None,
+                        "storage_options": {},
+                    },
                 ),
-                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.nullable_schema(
-                                    core_schema.str_schema(),
-                                ),
-                                default=None,
-                            ),
-                            required=False,
+            ]
+        )
+
+        object_schema = 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.nullable_schema(
+                            core_schema.str_schema(),
                         ),
-                        "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,
+                        default=None,
+                    ),
+                    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(),
                         ),
-                    },
-                    extra_behavior="forbid",
+                        default_factory=dict,
+                    ),
+                    required=False,
                 ),
+            },
+            extra_behavior="forbid",
+        )
+
+        deserialization_schema = core_schema.chain_schema(
+            [
+                core_schema.union_schema([str_schema, object_schema]),
                 core_schema.no_info_plain_validator_function(
                     lambda dct: cls(
                         dct.pop("path"),

upath/tests/test_pydantic.py

@@ -111,3 +111,8 @@ def test_dump_non_serializable_json():
         pydantic.TypeAdapter(UPath).dump_python(
             UPath("https://www.example.com", get_client=get_client), mode="json"
         )
+
+
+def test_json_schema():
+    ta = pydantic.TypeAdapter(UPath)
+    ta.json_schema()