From 24a10f7347949cfe9b4392f370913c4cadd5b437 Mon Sep 17 00:00:00 2001 From: SimonBoothroyd Date: Sun, 3 Nov 2024 08:01:58 -0500 Subject: [PATCH] fix: Don't process class aliases, as real classes are processed at some point anyway MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Processing the alias skips the real class later, which prevents the extension from seeing model fields. Issue-8: https://github.com/mkdocstrings/griffe-pydantic/issues/8 PR-7: https://github.com/mkdocstrings/griffe-pydantic/pull/7 Co-authored-by: Timothée Mazzucotelli --- src/griffe_pydantic/static.py | 4 +++- tests/test_extension.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/griffe_pydantic/static.py b/src/griffe_pydantic/static.py index d7fd046..70d6c91 100644 --- a/src/griffe_pydantic/static.py +++ b/src/griffe_pydantic/static.py @@ -162,7 +162,9 @@ def process_module( processed.add(mod.canonical_path) for cls in mod.classes.values(): - process_class(cls, processed=processed, schema=schema) + # Don't process aliases, real classes will be processed at some point anyway. + if not cls.is_alias: + process_class(cls, processed=processed, schema=schema) for submodule in mod.modules.values(): process_module(submodule, processed=processed, schema=schema) diff --git a/tests/test_extension.py b/tests/test_extension.py index b29340e..847e256 100644 --- a/tests/test_extension.py +++ b/tests/test_extension.py @@ -70,3 +70,17 @@ def test_extension() -> None: schema = package.classes["ExampleModel"].extra["griffe_pydantic"]["schema"] assert schema.startswith('{\n "description"') + + +def test_imported_models() -> None: + """Test the extension with imported models.""" + with temporary_visited_package( + "package", + modules={ + "__init__.py": "from ._private import MyModel\n\n__all__ = ['MyModel']", + "_private.py": "from pydantic import BaseModel\n\nclass MyModel(BaseModel):\n field1: str\n '''Some field.'''\n", + }, + extensions=Extensions(PydanticExtension(schema=True)), + ) as package: + assert package["MyModel"].labels == {"pydantic-model"} + assert package["MyModel.field1"].labels == {"pydantic-field"}