Refactor attrs support into reusable code

pydoctor/astutils.py

@@ -2,18 +2,22 @@
 Various bits of reusable code related to L{ast.AST} node processing.
 """
 
+import enum
 import inspect
 import platform
 import sys
 from numbers import Number
-from typing import Iterator, Optional, List, Iterable, Sequence, TYPE_CHECKING, Tuple, Union
+from typing import Iterator, Optional, List, Iterable, Sequence, TYPE_CHECKING, Tuple, Type, TypeVar, Union
 from inspect import BoundArguments, Signature
 import ast
 
 from pydoctor import visitor
 
 if TYPE_CHECKING:
     from pydoctor import model
+    from typing import Protocol, Literal
+else:
+    Protocol = Literal = object
 
 # AST visitors
 
@@ -402,4 +406,96 @@ def extract_docstring(node: ast.Str) -> Tuple[int, str]:
         - The docstring to be parsed, cleaned by L{inspect.cleandoc}.
     """
     lineno = extract_docstring_linenum(node)
-    return lineno, inspect.cleandoc(node.s)
+    return lineno, inspect.cleandoc(node.s)
+
+def safe_bind_args(sig:Signature, call: ast.AST, ctx: 'model.Module') -> Optional[inspect.BoundArguments]:
+    """
+    Binds the arguments of a function call to that function's signature.
+
+    When L{bind_args} raises a L{TypeError}, it reports a warning and returns C{None}. 
+    """
+    if not isinstance(call, ast.Call):
+        return None
+    try:
+        return bind_args(sig, call)
+    except TypeError as ex:
+        message = str(ex).replace("'", '"')
+        call_dottedname = node2dottedname(call.func)
+        callable_name = f"{'.'.join(call_dottedname)}()" if call_dottedname else 'callable'
+        ctx.report(
+            f"Invalid arguments for {callable_name}: {message}",
+            lineno_offset=call.lineno
+            )
+        return None
+
+class _V(enum.Enum): 
+    NoValue = enum.auto()
+_T =  TypeVar('_T', bound=object)
+def _get_literal_arg(args:BoundArguments, name:str, typecheck:Type[_T]) -> Union['Literal[_V.NoValue]', _T]:
+    """
+    Helper function for L{get_literal_arg}. 
+
+    If the value is not present in the arguments, returns L{_V.NoValue}.
+    @raises ValueError: If the passed value is not a literal or if it's not the right type.
+    """
+    expr = args.arguments.get(name)
+    if expr is None:
+        return _V.NoValue
+
+    try:
+        value = ast.literal_eval(expr)
+    except ValueError:
+        message = (
+            f'Unable to figure out value for {name!r} argument, maybe too complex'
+            ).replace("'", '"')
+        raise ValueError(message)
+
+    if not isinstance(value, typecheck):
+        message = (f'Value for {name!r} argument '
+            f'has type "{type(value).__name__}", expected {typecheck.__name__!r}'
+            ).replace("'", '"')
+        raise ValueError(message)
+
+    return value
+
+def get_literal_arg(args:BoundArguments, name:str, default:_T, 
+                          typecheck:Type[_T], lineno:int, module: 'model.Module') -> _T:
+    """
+    Retreive the literal value of an argument from the L{BoundArguments}. 
+    Only works with purely literal values (no C{Name} or C{Attribute}).
+
+    @param args: The L{BoundArguments} instance.
+    @param name: The name of the argument
+    @param default: The default value of the argument, this value is returned 
+        if the argument is not found.
+    @param typecheck: The type of the literal value this argument is expected to have.
+    @param lineno: The lineumber of the callsite, usd for error reporting.
+    @param module: Module that contains the call, used for error reporting.
+    @return: The value of the argument if we can infer it, otherwise returns
+        the default value.
+    """
+    try:
+        value = _get_literal_arg(args, name, typecheck)
+    except ValueError as e:
+        module.report(str(e), lineno_offset=lineno)
+        return default
+    if value is _V.NoValue:
+        # default value
+        return default
+    else:
+        return value
+
+class _HasDecoratorList(Protocol):
+    decorator_list:List[ast.expr]
+
+def iter_decorators(node:_HasDecoratorList, ctx: 'model.Documentable') -> Iterator[Tuple[Optional[str], ast.AST]]:
+    """
+    Utility function to iterate decorators.
+    """
+
+    for decnode in node.decorator_list:
+        namenode = decnode
+        if isinstance(namenode, ast.Call):
+            namenode = namenode.func
+        dottedname = node2fullname(namenode, ctx)
+        yield dottedname, decnode

pydoctor/extensions/_dataclass_like.py

@@ -0,0 +1,65 @@
+"""
+Dataclass-like libraries are all alike: 
+- They transform class variables into instance variable un certain conditions.
+- They autoamtically provides a constructor method without having to define __init__.
+
+More specifically
+"""
+import ast
+from abc import abstractmethod, ABC
+from typing import Optional, Union
+from pydoctor import astutils
+from pydoctor.model import Module, Attribute, Class, Documentable
+from pydoctor.extensions import ModuleVisitorExt, ClassMixin
+
+class DataclasLikeClass(ClassMixin):
+    isDataclassLike:bool = False
+
+class DataclassLikeVisitor(ModuleVisitorExt, ABC):
+
+    @abstractmethod
+    def isDataclassLike(self, cls:ast.ClassDef, mod:Module) -> bool:
+        """
+        Whether L{transformClassVar} method should be called for each class variables
+        in this class.
+        """
+
+    @abstractmethod
+    def transformClassVar(self, cls:Class, attr:Attribute, 
+                          annotation:Optional[ast.expr],
+                          value:Optional[ast.expr]) -> None:
+        """
+        Transform this class variable into a instance variable.
+        This method is left abstract because it might not be as simple as setting::
+            attr.kind = model.DocumentableKind.INSTANCE_VARIABLE
+        (but it also might be just that for the simpler cases)
+        """
+
+    def visit_ClassDef(self, node: ast.ClassDef) -> None:
+        cls = self.visitor.builder._stack[-1].contents.get(node.name)
+        if not isinstance(cls, Class):
+            return
+        assert isinstance(cls, DataclasLikeClass)
+        cls.isDataclassLike = self.isDataclassLike(node, cls.module)
+
+    def visit_Assign(self, node: Union[ast.Assign, ast.AnnAssign]) -> None:
+        current = self.visitor.builder.current
+
+        for dottedname in astutils.iterassign(node):
+            if dottedname and len(dottedname)==1:
+                # We consider single name assignment only
+                if not isinstance(current, Class):
+                    continue
+                assert isinstance(current, DataclasLikeClass)
+                if not current.isDataclassLike:
+                    continue
+                target, = dottedname
+                attr: Optional[Documentable] = current.contents.get(target)
+                if not isinstance(attr, Attribute):
+                    continue
+                if astutils.is_using_typing_classvar(attr.annotation, current):
+                    continue
+                annotation = node.annotation if isinstance(node, ast.AnnAssign) else None
+                self.transformClassVar(current, attr, annotation, node.value)
+
+    visit_AnnAssign = visit_Assign

pydoctor/extensions/attrs.py

@@ -6,9 +6,10 @@
 import ast
 import inspect
 
-from typing import Optional, Union
+from typing import Optional
 
 from pydoctor import astbuilder, model, astutils, extensions
+from pydoctor.extensions._dataclass_like import DataclasLikeClass, DataclassLikeVisitor
 
 import attr
 
@@ -18,80 +19,19 @@
 attrib_signature = inspect.signature(attr.ib)
 """Signature of the L{attr.ib} function for defining class attributes."""
 
-def uses_auto_attribs(call: ast.AST, module: model.Module) -> bool:
-    """Does the given L{attr.s()} decoration contain C{auto_attribs=True}?
-    @param call: AST of the call to L{attr.s()}.
-        This function will assume that L{attr.s()} is called without
-        verifying that.
-    @param module: Module that contains the call, used for error reporting.
-    @return: L{True} if L{True} is passed for C{auto_attribs},
-        L{False} in all other cases: if C{auto_attribs} is not passed,
-        if an explicit L{False} is passed or if an error was reported.
-    """
-    if not isinstance(call, ast.Call):
-        return False
-    if not astutils.node2fullname(call.func, module) in ('attr.s', 'attr.attrs', 'attr.attributes'):
-        return False
-    try:
-        args = astutils.bind_args(attrs_decorator_signature, call)
-    except TypeError as ex:
-        message = str(ex).replace("'", '"')
-        module.report(
-            f"Invalid arguments for attr.s(): {message}",
-            lineno_offset=call.lineno
-            )
-        return False
-
-    auto_attribs_expr = args.arguments.get('auto_attribs')
-    if auto_attribs_expr is None:
-        return False
-
-    try:
-        value = ast.literal_eval(auto_attribs_expr)
-    except ValueError:
-        module.report(
-            'Unable to figure out value for "auto_attribs" argument '
-            'to attr.s(), maybe too complex',
-            lineno_offset=call.lineno
-            )
-        return False
-
-    if not isinstance(value, bool):
-        module.report(
-            f'Value for "auto_attribs" argument to attr.s() '
-            f'has type "{type(value).__name__}", expected "bool"',
-            lineno_offset=call.lineno
-            )
-        return False
-
-    return value
+def is_attrs_deco(deco: ast.AST, module: model.Module) -> bool:
+    if isinstance(deco, ast.Call):
+        deco = deco.func
+    return astutils.node2fullname(deco, module) in (
+        'attr.s', 'attr.attrs', 'attr.attributes')
 
 def is_attrib(expr: Optional[ast.expr], ctx: model.Documentable) -> bool:
     """Does this expression return an C{attr.ib}?"""
     return isinstance(expr, ast.Call) and astutils.node2fullname(expr.func, ctx) in (
         'attr.ib', 'attr.attrib', 'attr.attr'
         )
 
-def attrib_args(expr: ast.expr, ctx: model.Documentable) -> Optional[inspect.BoundArguments]:
-    """Get the arguments passed to an C{attr.ib} definition.
-    @return: The arguments, or L{None} if C{expr} does not look like
-        an C{attr.ib} definition or the arguments passed to it are invalid.
-    """
-    if isinstance(expr, ast.Call) and astutils.node2fullname(expr.func, ctx) in (
-            'attr.ib', 'attr.attrib', 'attr.attr'
-            ):
-        try:
-            return astutils.bind_args(attrib_signature, expr)
-        except TypeError as ex:
-            message = str(ex).replace("'", '"')
-            ctx.module.report(
-                f"Invalid arguments for attr.ib(): {message}",
-                lineno_offset=expr.lineno
-                )
-    return None
-
 def annotation_from_attrib(
-        self: astbuilder.ModuleVistor,
         expr: ast.expr,
         ctx: model.Documentable
         ) -> Optional[ast.expr]:
@@ -101,7 +41,9 @@ def annotation_from_attrib(
     @return: A type annotation, or None if the expression is not
                 an C{attr.ib} definition or contains no type information.
     """
-    args = attrib_args(expr, ctx)
+    args = None
+    if is_attrib(expr, ctx):
+        args = astutils.safe_bind_args(attrib_signature, expr, ctx.module)
     if args is not None:
         typ = args.arguments.get('type')
         if typ is not None:
@@ -111,63 +53,51 @@ def annotation_from_attrib(
             return astbuilder._infer_type(default)
     return None
 
-class ModuleVisitor(extensions.ModuleVisitorExt):
+class ModuleVisitor(DataclassLikeVisitor):
 
     def visit_ClassDef(self, node:ast.ClassDef) -> None:
         """
         Called when a class definition is visited.
         """
-        cls = self.visitor.builder.current
-        if not isinstance(cls, model.Class) or cls.name!=node.name:
-            return
+        super().visit_ClassDef(node)
 
-        assert isinstance(cls, AttrsClass)
-        cls.auto_attribs = any(uses_auto_attribs(decnode, cls.module) for decnode in node.decorator_list)
+        cls = self.visitor.builder._stack[-1].contents.get(node.name)
+        if not isinstance(cls, AttrsClass) or not cls.isDataclassLike:
+            return
 
-    def _handleAttrsAssignmentInClass(self, target:str, node: Union[ast.Assign, ast.AnnAssign]) -> None:
-        cls = self.visitor.builder.current
-        assert isinstance(cls, AttrsClass)
+        for name, decnode in astutils.iter_decorators(node, cls):
+            if not name in ('attr.s', 'attr.attrs', 'attr.attributes'):
+                continue
 
-        attr: Optional[model.Documentable] = cls.contents.get(target)
-        if attr is None:
-            return
-        if not isinstance(attr, model.Attribute):
-            return
+            attrs_args = astutils.safe_bind_args(attrs_decorator_signature, decnode, cls.module)
+            if attrs_args:
 
-        annotation = node.annotation if isinstance(node, ast.AnnAssign) else None
-
-        if is_attrib(node.value, cls) or (
-               cls.auto_attribs and \
-               annotation is not None and \
-               not astutils.is_using_typing_classvar(annotation, cls)):
-
+                cls.auto_attribs = astutils.get_literal_arg(
+                    name='auto_attribs', default=False, typecheck=bool,
+                    args=attrs_args, lineno=decnode.lineno, module=cls.module)
+
+            break
+
+    def transformClassVar(self, cls: model.Class, 
+                          attr: model.Attribute, 
+                          annotation:Optional[ast.expr],
+                          value:Optional[ast.expr]) -> None:
+        assert isinstance(cls, AttrsClass)
+        if is_attrib(value, cls) or (cls.auto_attribs and annotation is not None):
             attr.kind = model.DocumentableKind.INSTANCE_VARIABLE
-            if annotation is None and node.value is not None:
-                attr.annotation = annotation_from_attrib(self.visitor, node.value, cls)
-
-    def _handleAttrsAssignment(self, node: Union[ast.Assign, ast.AnnAssign]) -> None:
-        for dottedname in astutils.iterassign(node):
-            if dottedname and len(dottedname)==1:
-                # Here, we consider single name assignment only
-                current = self.visitor.builder.current
-                if isinstance(current, model.Class):
-                    self._handleAttrsAssignmentInClass(
-                        dottedname[0], node
-                    )
-
-    def visit_Assign(self, node: Union[ast.Assign, ast.AnnAssign]) -> None:
-        self._handleAttrsAssignment(node)
-    visit_AnnAssign = visit_Assign
-
-class AttrsClass(extensions.ClassMixin, model.Class):
+            if annotation is None and value is not None:
+                attr.annotation = annotation_from_attrib(value, cls)
 
-    def setup(self) -> None:
-        super().setup()
-        self.auto_attribs: bool = False
-        """
-        L{True} if this class uses the C{auto_attribs} feature of the L{attrs}
-        library to automatically convert annotated fields into attributes.
-        """
+    def isDataclassLike(self, cls:ast.ClassDef, mod:model.Module) -> bool:
+        return any(is_attrs_deco(dec, mod) for dec in cls.decorator_list)
+
+class AttrsClass(DataclasLikeClass, model.Class):
+
+    auto_attribs: bool = False
+    """
+    L{True} if this class uses the C{auto_attribs} feature of the L{attrs}
+    library to automatically convert annotated fields into attributes.
+    """
 
 def setup_pydoctor_extension(r:extensions.ExtRegistrar) -> None:
     r.register_astbuilder_visitor(ModuleVisitor)