From 2439409c000ec9afee476685ea0a04149f5564bf Mon Sep 17 00:00:00 2001 From: Travis Prescott Date: Tue, 7 Mar 2023 09:25:21 -0800 Subject: [PATCH] Display class decorators. Fixes #4860 --- .../api-stub-generator/CHANGELOG.md | 1 + .../apistub/nodes/_class_node.py | 15 +++++++++++++++ .../apistub/nodes/_data_class_node.py | 1 + .../tests/class_parsing_test.py | 18 ++++++++++++++++++ .../apistubgentest/models/__init__.py | 2 ++ .../apistubgentest/models/_models.py | 17 +++++++++++++++++ 6 files changed, 54 insertions(+) diff --git a/packages/python-packages/api-stub-generator/CHANGELOG.md b/packages/python-packages/api-stub-generator/CHANGELOG.md index 1ed7770e371..1dc5b63afb6 100644 --- a/packages/python-packages/api-stub-generator/CHANGELOG.md +++ b/packages/python-packages/api-stub-generator/CHANGELOG.md @@ -3,6 +3,7 @@ ## Version 0.3.7 (Unreleased) Fix incorrect type annotation. Update to follow best practices for accessing '__annotations__'. +Fixed issue where class decorators were not displayed. ## Version 0.3.6 (2022-10-27) Suppressed unwanted base class methods in DPG libraries. diff --git a/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py b/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py index 68167e6bdf0..1189df53c42 100644 --- a/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py +++ b/packages/python-packages/api-stub-generator/apistub/nodes/_class_node.py @@ -135,6 +135,14 @@ def _handle_class_variable(self, child_obj, name, *, type_string=None, value=Non ) ) + def _parse_decorators_from_class(self, class_obj): + try: + class_node = astroid.parse(inspect.getsource(class_obj)).body[0] + class_decorators = class_node.decorators.nodes + self.decorators = [f"@{x.as_string(preserve_quotes=True)}" for x in class_decorators] + except: + self.decorators = [] + def _parse_functions_from_class(self, class_obj) -> List[astroid.FunctionDef]: try: class_node = astroid.parse(inspect.getsource(class_obj)).body[0] @@ -179,6 +187,8 @@ def _inspect(self): is_typeddict = hasattr(self.obj, "__required_keys__") or hasattr(self.obj, "__optional_keys__") + self._parse_decorators_from_class(self.obj) + # find members in node # enums with duplicate values are screened out by "getmembers" so # we must rely on __members__ instead. @@ -297,6 +307,11 @@ def generate_tokens(self, apiview): """ logging.info(f"Processing class {self.namespace_id}") # Generate class name line + for decorator in self.decorators: + apiview.add_whitespace() + apiview.add_keyword(decorator) + apiview.add_newline() + apiview.add_whitespace() apiview.add_line_marker(self.namespace_id) apiview.add_keyword("class", False, True) diff --git a/packages/python-packages/api-stub-generator/apistub/nodes/_data_class_node.py b/packages/python-packages/api-stub-generator/apistub/nodes/_data_class_node.py index e470d7ba2bd..a207baefa68 100644 --- a/packages/python-packages/api-stub-generator/apistub/nodes/_data_class_node.py +++ b/packages/python-packages/api-stub-generator/apistub/nodes/_data_class_node.py @@ -14,6 +14,7 @@ class DataClassNode(ClassNode): def __init__(self, *, name, namespace, parent_node, obj, pkg_root_namespace): super().__init__(name=name, namespace=namespace, parent_node=parent_node, obj=obj, pkg_root_namespace=pkg_root_namespace) + self.decorators = [x for x in self.decorators if not x.startswith("@dataclass")] # explicitly set synthesized __init__ return type to None to fix test flakiness for child in self.child_nodes: if child.display_name == "__init__": diff --git a/packages/python-packages/api-stub-generator/tests/class_parsing_test.py b/packages/python-packages/api-stub-generator/tests/class_parsing_test.py index 2eae6d61591..b96fddcc38b 100644 --- a/packages/python-packages/api-stub-generator/tests/class_parsing_test.py +++ b/packages/python-packages/api-stub-generator/tests/class_parsing_test.py @@ -9,6 +9,7 @@ from apistubgentest.models import ( AliasNewType, AliasUnion, + ClassWithDecorators, FakeTypedDict, FakeObject, GenericStack, @@ -39,6 +40,23 @@ class TestClassParsing: pkg_namespace = "apistubgentest.models" + def test_class_with_decorators(self): + obj = ClassWithDecorators + class_node = ClassNode(name=obj.__name__, namespace=obj.__name__, parent_node=None, obj=obj, pkg_root_namespace=self.pkg_namespace) + actuals = _render_lines(_tokenize(class_node)) + expected = [ + "@add_id", + "class ClassWithDecorators:", + "", + "def __init__(", + "self, ", + "id, ", + "*args, ", + "**kwargs", + ")", + ] + _check_all(actuals, expected, obj) + def test_typed_dict_class(self): obj = FakeTypedDict class_node = ClassNode(name=obj.__name__, namespace=obj.__name__, parent_node=None, obj=obj, pkg_root_namespace=self.pkg_namespace) diff --git a/packages/python-packages/apistubgentest/apistubgentest/models/__init__.py b/packages/python-packages/apistubgentest/apistubgentest/models/__init__.py index 0099052fc71..32df3cef3dc 100644 --- a/packages/python-packages/apistubgentest/apistubgentest/models/__init__.py +++ b/packages/python-packages/apistubgentest/apistubgentest/models/__init__.py @@ -10,6 +10,7 @@ from ._models import ( AliasNewType, AliasUnion, + ClassWithDecorators, DocstringClass, FakeError, FakeObject, @@ -39,6 +40,7 @@ __all__ = ( "AliasNewType", "AliasUnion", + "ClassWithDecorators", "DataClassSimple", "DataClassWithFields", "DataClassDynamic", diff --git a/packages/python-packages/apistubgentest/apistubgentest/models/_models.py b/packages/python-packages/apistubgentest/apistubgentest/models/_models.py index 214aee5bbe2..74ea0e5e99d 100644 --- a/packages/python-packages/apistubgentest/apistubgentest/models/_models.py +++ b/packages/python-packages/apistubgentest/apistubgentest/models/_models.py @@ -31,6 +31,23 @@ def wrapper(*args, **kwargs): return wrapper return decorator +def get_id(self): + return self.__id + +def add_id(cls): + cls_init = cls.__init__ + + def __init__(self, id, *args, **kwargs): + self.__id = id + self.get_id = get_id + cls_init(self, *args, **kwargs) + + cls.__init__ = __init__ + return cls + +@add_id +class ClassWithDecorators: + pass class PublicCaseInsensitiveEnumMeta(EnumMeta): def __getitem__(self, name: str):