Skip to content

Commit

Permalink
Do not add imports if we added no type info in ATAV
Browse files Browse the repository at this point in the history
In ApplyTypeAnnotationsVisitor, there are edge cases where we
might have changed the module imports even though we never wound
up applying any type annotations.

This will become even more common if we support adding
`from __future__ import annotations`, which I would like to do
soon.

To handle this, we can simply return the original tree from
`transform_module_impl` (discarding any changes from either
`self` or `AddImportsVisitor`) whenever there are no changes
in `self.annotation_counts`.

I updated the no-annotations-changed test to reflect this:
```
> python -m unittest libcst.codemod.visitors.tests.test_apply_type_annotations.TestApplyAnnotationsVisitor
...............................................
----------------------------------------------------------------------
Ran 47 tests in 2.312s

OK
```
  • Loading branch information
stroxler committed Oct 28, 2021
1 parent 3743c70 commit 87625d0
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 8 deletions.
10 changes: 8 additions & 2 deletions libcst/codemod/visitors/_apply_type_annotations.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,7 @@ class AnnotationCounts:
return_annotations: int = 0
classes_added: int = 0

def applied_changes(self):
def any_changes(self):
return (
self.global_annotations
+ self.attribute_annotations
Expand Down Expand Up @@ -385,7 +385,13 @@ def transform_module_impl(self, tree: cst.Module) -> cst.Module:
self.annotations.class_definitions.update(visitor.class_definitions)

tree_with_imports = AddImportsVisitor(self.context).transform_module(tree)
return tree_with_imports.visit(self)
tree_with_changes = tree_with_imports.visit(self)

# don't modify the imports if we didn't actually add any type information
if self.annotation_counts.any_changes():
return tree_with_changes
else:
return tree

# smart constructors: all applied annotations happen via one of these

Expand Down
8 changes: 2 additions & 6 deletions libcst/codemod/visitors/tests/test_apply_type_annotations.py
Original file line number Diff line number Diff line change
Expand Up @@ -1012,11 +1012,7 @@ class C:
def __init__(self):
self.attr_will_not_be_found = None
""",
# TODO: use the annotation counts to avoid adding
# the import in this case.
"""
from bar import X
class C:
def __init__(self):
self.attr_will_not_be_found = None
Expand All @@ -1032,7 +1028,7 @@ def test_count_annotations(
before: str,
after: str,
annotation_counts: AnnotationCounts,
applied_changes: False,
any_changes: False,
):
stub = self.make_fixture_data(stub)
before = self.make_fixture_data(before)
Expand All @@ -1048,4 +1044,4 @@ def test_count_annotations(

self.assertEqual(after, output_code)
self.assertEqual(str(annotation_counts), str(visitor.annotation_counts))
self.assertEqual(applied_changes, visitor.annotation_counts.applied_changes())
self.assertEqual(any_changes, visitor.annotation_counts.any_changes())

0 comments on commit 87625d0

Please sign in to comment.