Skip to content

Commit

Permalink
Add type annotations to pyreverse dot files
Browse files Browse the repository at this point in the history
  • Loading branch information
mbyrnepr2 committed Jun 9, 2021
1 parent 3ca2b25 commit caed9b0
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 4 deletions.
4 changes: 4 additions & 0 deletions ChangeLog
Expand Up @@ -9,6 +9,10 @@ Release date: TBA
..
Put new features and bugfixes here and also in 'doc/whatsnew/2.9.rst'

* Add type annotations to pyreverse dot files

Closes #1548

* Added ``deprecated-decorator``: Emitted when deprecated decorator is used.

Closes #4429
Expand Down
2 changes: 2 additions & 0 deletions doc/whatsnew/2.9.rst
Expand Up @@ -40,6 +40,8 @@ New checkers
Other Changes
=============

* Add type annotations to pyreverse dot files

* Pylint's tags are now the standard form ``vX.Y.Z`` and not ``pylint-X.Y.Z`` anymore.

* Fix false-positive ``too-many-ancestors`` when inheriting from builtin classes,
Expand Down
2 changes: 1 addition & 1 deletion pylint/pyreverse/diagrams.py
Expand Up @@ -122,7 +122,7 @@ def class_names(self, nodes):
if isinstance(node, astroid.Instance):
node = node._proxied
if (
isinstance(node, astroid.ClassDef)
isinstance(node, (astroid.ClassDef, astroid.Name))
and hasattr(node, "name")
and not self.has_node(node)
):
Expand Down
16 changes: 14 additions & 2 deletions pylint/pyreverse/inspector.py
Expand Up @@ -218,8 +218,10 @@ def visit_assignname(self, node):
self.visit_module(frame)

current = frame.locals_type[node.name]
values = set(node.infer())
ann = getattr(node.parent, "annotation", None)
values = {ann} if ann else set(node.infer())
frame.locals_type[node.name] = list(set(current) | values)

except astroid.InferenceError:
pass

Expand All @@ -229,8 +231,18 @@ def handle_assignattr_type(node, parent):
handle instance_attrs_type
"""

ann = None
if isinstance(node.parent.value, astroid.Name):
try:
idx = list(node.parent.parent.locals).index(node.parent.value.name)
ann = node.parent.parent.args.annotations[idx]
except IndexError:
pass
elif hasattr(node.parent, "annotation"):
ann = node.parent.annotation
try:
values = set(node.infer())
values = {ann} if ann else set(node.infer())
current = set(parent.instance_attrs_type[node.attrname])
parent.instance_attrs_type[node.attrname] = list(current | values)
except astroid.InferenceError:
Expand Down
11 changes: 10 additions & 1 deletion pylint/pyreverse/writer.py
Expand Up @@ -134,11 +134,20 @@ def get_values(self, obj):
if not self.config.only_classnames:
label = r"{}|{}\l|".format(label, r"\l".join(obj.attrs))
for func in obj.methods:
return_type = f": {func.returns.name}" if func.returns else ""

if func.args.args:
args = [arg.name for arg in func.args.args if arg.name != "self"]
else:
args = []
label = r"{}{}({})\l".format(label, func.name, ", ".join(args))

annotations = [a.name if a else None for a in func.args.annotations[1:]]
args = ", ".join(
f"{arg}: {ann}" if ann else f"{arg}"
for arg, ann in zip(args, annotations)
)

label = fr"{label}{func.name}({args}){return_type}\l"
label = "{%s}" % label
if is_exception(obj.node):
return dict(fontcolor="red", label=label, shape="record")
Expand Down

0 comments on commit caed9b0

Please sign in to comment.