Added subscript visitor method.

spcl · Sep 14, 2023 · 0bbe5c4 · 0bbe5c4
1 parent 36010fa
commit 0bbe5c4
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/dace/symbolic.py b/dace/symbolic.py
@@ -1027,6 +1027,15 @@ def visit_IfExp(self, node):
                                   self.visit(node.orelse)],
                             keywords=[])
         return ast.copy_location(new_node, node)
+
+    def visit_Subscript(self, node):
+        if isinstance(node.value, ast.Attribute):
+            attr = ast.Subscript(value=ast.Name(id=node.value.attr, ctx=ast.Load()), slice=node.slice, ctx=ast.Load())
+            new_node = ast.Call(func=ast.Name(id='Attr', ctx=ast.Load),
+                                args=[self.visit(node.value.value), self.visit(attr)],
+                                keywords=[])
+            return ast.copy_location(new_node, node)
+        return self.generic_visit(node)
 
     def visit_Attribute(self, node):
         new_node = ast.Call(func=ast.Name(id='Attr', ctx=ast.Load),