feat: pydantic v2 support

Author: khalo-sa

.gitignore

@@ -107,6 +107,7 @@ celerybeat.pid
 # Environments
 .env
 .venv
+.venv-v1
 env/
 venv/
 ENV/
@@ -143,3 +144,6 @@ cython_debug/
 
 # VS Code config
 .vscode
+
+# test outputs
+output_debug.ts

pydantic2ts/cli/script.py

@@ -7,12 +7,13 @@
 import shutil
 import sys
 from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
 from tempfile import mkdtemp
 from types import ModuleType
 from typing import Any, Dict, List, Tuple, Type
 from uuid import uuid4
 
-from pydantic import BaseModel, Extra, create_model
+from pydantic import VERSION, BaseModel, Extra, create_model
 
 try:
     from pydantic.generics import GenericModel
@@ -22,6 +23,11 @@
 logger = logging.getLogger("pydantic2ts")
 
 
+DEBUG = os.environ.get("DEBUG", False)
+
+V2 = True if VERSION.startswith("2") else False
+
+
 def import_module(path: str) -> ModuleType:
     """
     Helper which allows modules to be specified by either dotted path notation or by filepath.
@@ -61,12 +67,15 @@ def is_concrete_pydantic_model(obj) -> bool:
     Return true if an object is a concrete subclass of pydantic's BaseModel.
     'concrete' meaning that it's not a GenericModel.
     """
+    generic_metadata = getattr(obj, "__pydantic_generic_metadata__", None)
     if not inspect.isclass(obj):
         return False
     elif obj is BaseModel:
         return False
-    elif GenericModel and issubclass(obj, GenericModel):
+    elif not V2 and GenericModel and issubclass(obj, GenericModel):
         return bool(obj.__concrete__)
+    elif V2 and generic_metadata:
+        return not bool(generic_metadata["parameters"])
     else:
         return issubclass(obj, BaseModel)
 
@@ -141,7 +150,7 @@ def clean_schema(schema: Dict[str, Any]) -> None:
         del schema["description"]
 
 
-def generate_json_schema(models: List[Type[BaseModel]]) -> str:
+def generate_json_schema_v1(models: List[Type[BaseModel]]) -> str:
     """
     Create a top-level '_Master_' model with references to each of the actual models.
     Generate the schema for this model, which will include the schemas for all the
@@ -178,6 +187,43 @@ def generate_json_schema(models: List[Type[BaseModel]]) -> str:
                 m.Config.extra = x
 
 
+def generate_json_schema_v2(models: List[Type[BaseModel]]) -> str:
+    """
+    Create a top-level '_Master_' model with references to each of the actual models.
+    Generate the schema for this model, which will include the schemas for all the
+    nested models. Then clean up the schema.
+
+    One weird thing we do is we temporarily override the 'extra' setting in models,
+    changing it to 'forbid' UNLESS it was explicitly set to 'allow'. This prevents
+    '[k: string]: any' from being added to every interface. This change is reverted
+    once the schema has been generated.
+    """
+    model_extras = [m.model_config.get("extra") for m in models]
+
+    try:
+        for m in models:
+            if m.model_config.get("extra") != Extra.allow:
+                m.model_config["extra"] = Extra.forbid
+
+        master_model: BaseModel = create_model(
+            "_Master_", **{m.__name__: (m, ...) for m in models}
+        )
+        master_model.model_config["extra"] = Extra.forbid
+        master_model.model_config["json_schema_extra"] = staticmethod(clean_schema)
+
+        schema: dict = master_model.model_json_schema()
+
+        for d in schema.get("$defs", {}).values():
+            clean_schema(d)
+
+        return json.dumps(schema, indent=2)
+
+    finally:
+        for m, x in zip(models, model_extras):
+            if x is not None:
+                m.model_config["extra"] = x
+
+
 def generate_typescript_defs(
     module: str, output: str, exclude: Tuple[str] = (), json2ts_cmd: str = "json2ts"
 ) -> None:
@@ -205,13 +251,18 @@ def generate_typescript_defs(
 
     logger.info("Generating JSON schema from pydantic models...")
 
-    schema = generate_json_schema(models)
+    schema = generate_json_schema_v2(models) if V2 else generate_json_schema_v1(models)
+
     schema_dir = mkdtemp()
     schema_file_path = os.path.join(schema_dir, "schema.json")
 
     with open(schema_file_path, "w") as f:
         f.write(schema)
 
+    if DEBUG:
+        with open(Path(output).parent / "schema.json", "w") as f:
+            f.write(schema)
+
     logger.info("Converting JSON schema to typescript definitions...")
 
     json2ts_exit_code = os.system(

tests/expected_results/excluding_models/input.py renamed to tests/expected_results/excluding_models/v1/input.py

@@ -0,0 +1,12 @@
+/* tslint:disable */
+/* eslint-disable */
+/**
+/* This file was automatically generated from pydantic models by running pydantic2ts.
+/* Do not modify it by hand - just update the pydantic models and then re-run the script
+*/
+
+export interface Profile {
+  username: string;
+  age: number | null;
+  hobbies: string[];
+}

tests/expected_results/excluding_models/output.ts renamed to tests/expected_results/excluding_models/v1/output.ts

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Generic, TypeVar, Optional, List, Type, cast, Union
+from typing import Generic, List, Optional, Type, TypeVar, cast
 
 from pydantic import BaseModel
 from pydantic.generics import GenericModel

tests/expected_results/single_module/input.py renamed to tests/expected_results/excluding_models/v2/input.py

@@ -0,0 +1,52 @@
+from datetime import datetime
+from typing import Generic, List, Optional, Type, TypeVar, cast
+
+from pydantic import BaseModel
+
+T = TypeVar("T")
+
+
+class Error(BaseModel):
+    code: int
+    message: str
+
+
+class ApiResponse(BaseModel, Generic[T]):
+    data: Optional[T]
+    error: Optional[Error]
+
+
+def create_response_type(data_type: T, name: str) -> "Type[ApiResponse[T]]":
+    """
+    Create a concrete implementation of ApiResponse and then applies the specified name.
+    This is necessary because the name automatically generated by __concrete_name__ is
+    really ugly, it just doesn't look good.
+    """
+    t = ApiResponse[data_type]
+    t.__name__ = name
+    t.__qualname__ = name
+    return cast(Type[ApiResponse[T]], t)
+
+
+class User(BaseModel):
+    name: str
+    email: str
+
+
+class UserProfile(User):
+    joined: datetime
+    last_active: datetime
+    age: int
+
+
+class Article(BaseModel):
+    author: User
+    content: str
+    published: datetime
+
+
+ListUsersResponse = create_response_type(List[User], "ListUsersResponse")
+
+ListArticlesResponse = create_response_type(List[Article], "ListArticlesResponse")
+
+UserProfileResponse = create_response_type(UserProfile, "UserProfileResponse")

tests/expected_results/excluding_models/v2/output.ts

@@ -0,0 +1,39 @@
+/* tslint:disable */
+/* eslint-disable */
+/**
+/* This file was automatically generated from pydantic models by running pydantic2ts.
+/* Do not modify it by hand - just update the pydantic models and then re-run the script
+*/
+
+export interface Article {
+  author: User;
+  content: string;
+  published: string;
+}
+export interface User {
+  name: string;
+  email: string;
+}
+export interface Error {
+  code: number;
+  message: string;
+}
+export interface ListArticlesResponse {
+  data: Article[] | null;
+  error: Error | null;
+}
+export interface ListUsersResponse {
+  data: User[] | null;
+  error: Error | null;
+}
+export interface UserProfile {
+  name: string;
+  email: string;
+  joined: string;
+  last_active: string;
+  age: number;
+}
+export interface UserProfileResponse {
+  data: UserProfile | null;
+  error: Error | null;
+}