Remove ast.Index as well

func_adl/ast/function_simplifier.py

@@ -433,14 +433,14 @@ def visit_Call(self, call_node):
         else:
             return FuncADLNodeTransformer.visit_Call(self, call_node)
 
-    def visit_Subscript_Tuple(self, v: ast.Tuple, s: Union[ast.Constant, ast.Index]):
+    def visit_Subscript_Tuple(self, v: ast.Tuple, s: ast.Constant):
         """
         (t1, t2, t3...)[1] => t2
 
         Only works if index is a number
         """
         # Get the value out - this is due to supporting python 3.7-3.9
-        n = _get_value_from_index(s)
+        n = s.value
         if n is None:
             return ast.Subscript(v, s, ast.Load())  # type: ignore
         assert isinstance(n, int), "Programming error: index is not an integer in tuple subscript"
@@ -452,13 +452,13 @@ def visit_Subscript_Tuple(self, v: ast.Tuple, s: Union[ast.Constant, ast.Index])
 
         return v.elts[n]
 
-    def visit_Subscript_List(self, v: ast.List, s: Union[ast.Constant, ast.Index]):
+    def visit_Subscript_List(self, v: ast.List, s: ast.Constant):
         """
         [t1, t2, t3...][1] => t2
 
         Only works if index is a number
         """
-        n = _get_value_from_index(s)
+        n = s.value
         if n is None:
             return ast.Subscript(v, s, ast.Load())  # type: ignore
         if n >= len(v.elts):
@@ -469,19 +469,19 @@ def visit_Subscript_List(self, v: ast.List, s: Union[ast.Constant, ast.Index]):
 
         return v.elts[n]
 
-    def visit_Subscript_Dict(self, v: ast.Dict, s: Union[ast.Constant, ast.Index]):
+    def visit_Subscript_Dict(self, v: ast.Dict, s: ast.Constant):
         """
         {t1, t2, t3...}[1] => t2
         """
-        sub = _get_value_from_index(s)
+        sub = s.value
         assert isinstance(sub, (str, int))
         return self.visit_Subscript_Dict_with_value(v, sub)
 
     def visit_Subscript_Dict_with_value(self, v: ast.Dict, s: Union[str, int]):
         "Do the lookup for the dict"
         for index, value in enumerate(v.keys):
             assert isinstance(value, ast.Constant)
-            if _get_value_from_index(value) == s:
+            if value.value == s:
                 return v.values[index]
 
         return ast.Subscript(v, s, ast.Load())  # type: ignore
@@ -561,23 +561,3 @@ def visit_Attribute(self, node):
             return self.visit_Subscript_Dict_with_value(visited_value, node.attr)
 
         return ast.Attribute(value=visited_value, attr=node.attr, ctx=ast.Load())
-
-
-def _get_value_from_index(arg: Union[ast.Constant, ast.Index]) -> Optional[int]:
-    """Deal with 3.7, and 3.8 differences in how indexing for list and tuple
-    subscripts is handled.
-
-    Args:
-        arg (Union[ast.Constant, ast.Index]): Input ast to extract an index from.
-                                                       Hopefully.
-    """
-
-    def extract(a: ast.Constant) -> Optional[int]:
-        if isinstance(a, ast.Constant):
-            return a.value
-        return None
-
-    if isinstance(arg, ast.Index):
-        return extract(arg.value)  # type: ignore
-    else:
-        return extract(arg)

func_adl/type_based_replacement.py

@@ -755,10 +755,7 @@ def process_parameterized_method_call(
                 )
 
             # Get the parameters from the subscript
-            if isinstance(slice, ast.Index):
-                parameters = ast.literal_eval(slice.value)  # type: ignore
-            else:
-                parameters = ast.literal_eval(slice)
+            parameters = ast.literal_eval(slice)
 
             # rebuild the expression, removing the slice operation and turning this into a
             # "normal" call.