InferredAnnAssignRewriter: Dont redeclare type

Also dont add annotation to subscript assignments. Closes py2many#360
MiguelMarcelino · Jul 4, 2021 · da9845b · da9845b
1 parent d0d30b9
commit da9845b
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 12 deletions.
diff --git a/py2many/rewriters.py b/py2many/rewriters.py
@@ -11,7 +11,9 @@
 
 class InferredAnnAssignRewriter(ast.NodeTransformer):
     def visit_Assign(self, node):
-        target = node.targets[0]
+        target = node.targets[0]  # Assumes all targets have same annotation
+        if isinstance(target, ast.Subscript):
+            return node
         annotation = getattr(target, "annotation", False)
         if not annotation:
             return node
@@ -23,8 +25,21 @@ def visit_Assign(self, node):
 
         assigns = []
         for assign_target in node.targets:
-            print(assign_target.__class__)
-            # assert False
+            definition = node.scopes.find(get_id(assign_target))
+            if definition is not assign_target:
+                previous_type = get_inferred_type(definition)
+                if get_id(previous_type) == get_id(annotation):
+                    if len(node.targets) == 1:
+                        return node
+                    else:
+                        new_node = ast.Assign(
+                            targets=[assign_target],
+                            value=node.value,
+                            lineno=node.lineno,
+                            col_offset=col_offset,
+                        )
+                        assigns.append(new_node)
+                        continue
             new_node = ast.AnnAssign(
                 target=assign_target,
                 value=node.value,

diff --git a/tests/expected/binit.py b/tests/expected/binit.py
@@ -11,9 +11,9 @@ def bisect_right(data: List[int], item: int) -> int:
     while low < high:
         middle: int = int((low + high) / 2)
         if item < data[middle]:
-            high: int = middle
+            high = middle
         else:
-            low: int = middle + 1
+            low = middle + 1
     return low
 
 

diff --git a/tests/expected/bubble_sort.py b/tests/expected/bubble_sort.py
@@ -12,8 +12,8 @@ def bubble_sort(seq: List[int]) -> List[int]:
             if seq[n] < seq[n - 1]:
                 if True:
                     (__tmp1, __tmp2) = (seq[n], seq[n - 1])
-                    seq[n - 1]: int = __tmp1
-                    seq[n]: int = __tmp2
+                    seq[n - 1] = __tmp1
+                    seq[n] = __tmp2
     return seq
 
 

diff --git a/tests/expected/comb_sort.py b/tests/expected/comb_sort.py
@@ -11,14 +11,14 @@ def comb_sort(seq: List[int]) -> List[int]:
     swap: bool = True
     while gap > 1 or swap:
         gap: int = max(1, floor(gap / 1.25))
-        swap: bool = False
+        swap = False
         for i in range(len(seq) - gap):
             if seq[i] > seq[i + gap]:
                 if True:
                     (__tmp1, __tmp2) = (seq[i + gap], seq[i])
-                    seq[i]: int = __tmp1
-                    seq[i + gap]: int = __tmp2
-                swap: bool = True
+                    seq[i] = __tmp1
+                    seq[i + gap] = __tmp2
+                swap = True
     return seq
 
 

diff --git a/tests/expected/rect.py b/tests/expected/rect.py
@@ -17,7 +17,7 @@ def is_square(self) -> bool:
 def show():
     r: Rectangle = Rectangle(height=1, length=1)
     assert r.is_square()
-    r: Rectangle = Rectangle(height=1, length=2)
+    r = Rectangle(height=1, length=2)
     assert not r.is_square()
     print(r.height)
     print(r.length)