added more Ast implementation classes

owasp-sbot · Jan 1, 2024 · 2896aef · 2896aef
1 parent 7a1a62f
commit 2896aef
Show file tree

Hide file tree

Showing 18 changed files with 258 additions and 33 deletions.
diff --git a/osbot_utils/utils/ast/Ast.py b/osbot_utils/utils/ast/Ast.py
@@ -20,4 +20,7 @@ def ast_module__from_source_code(self, source_code):
         if type(result) is ast.Module:
             return Ast_Module(result)
 
+    def parse(self, source_code):
+        return ast.parse(source_code)
+
 
diff --git a/osbot_utils/utils/ast/Ast_Arguments.py b/osbot_utils/utils/ast/Ast_Arguments.py
@@ -5,7 +5,7 @@ class Ast_Arguments(Ast_Node):
 
     def __init__(self, args):
         if type(args) is not ast.arguments:
-            raise Exception(f'Expected module to be of type ast.Module, got: {module}')
+            raise Exception(f'Expected module to be of type ast.arguments, got: {args}')
         super().__init__(args)
         self.args = args
 

diff --git a/osbot_utils/utils/ast/Ast_Assign.py b/osbot_utils/utils/ast/Ast_Assign.py
@@ -0,0 +1,11 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Assign(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Assign]'
+
+    def info(self):
+        return {'Ast_Assign': {'targets': self.targets(),
+                               'value'  : self.value()  }}
diff --git a/osbot_utils/utils/ast/Ast_Attribute.py b/osbot_utils/utils/ast/Ast_Attribute.py
@@ -0,0 +1,12 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Attribute(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Attribute]'
+
+    def info(self):
+        return {'Ast_Attribute': { 'attr'  : self.node.attr  ,
+                                   'ctx'   : self.ctx()}     ,
+                                   'value' : self.value()    }
diff --git a/osbot_utils/utils/ast/Ast_Call.py b/osbot_utils/utils/ast/Ast_Call.py
@@ -0,0 +1,12 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Call(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Call]'
+
+    def info(self):
+        return {'Ast_Call': { 'args'  : self.args()      ,
+                              'func'  : self.func()      ,
+                              'keywords': self.keywords()}}
diff --git a/osbot_utils/utils/ast/Ast_Class_Def.py b/osbot_utils/utils/ast/Ast_Class_Def.py
@@ -0,0 +1,12 @@
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Class_Def(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Class_Def]'
+
+    def info(self):
+        #self.print()
+        return {'Ast_Class_Def': {'bases': self.bases()  ,
+                                  'body' : self.body()   ,
+                                  'name' : self.node.name}}
diff --git a/osbot_utils/utils/ast/Ast_Dict.py b/osbot_utils/utils/ast/Ast_Dict.py
@@ -0,0 +1,11 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Dict(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Dict]'
+
+    def info(self):
+        return {'Ast_Dict': { 'keys'  : self.keys()  ,
+                              'values': self.values()}}
diff --git a/osbot_utils/utils/ast/Ast_Expr.py b/osbot_utils/utils/ast/Ast_Expr.py
@@ -0,0 +1,10 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Expr(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Expr]'
+
+    def info(self):
+        return {'Ast_Expr': { 'value'  : self.value()  }}
diff --git a/osbot_utils/utils/ast/Ast_Function_Def.py b/osbot_utils/utils/ast/Ast_Function_Def.py
@@ -12,10 +12,11 @@ def __repr__(self):
         return f'[Ast_Node][Ast_Function_Def]'
 
     def info(self):
-        info = obj_data(self.node)
-        #obj_info(self.node)
-        #return info
-        return {'Ast_Function_Def': {'args' : self.args().info() }}
+        return {'Ast_Function_Def': {'args'   : self.args().info(),
+                                     'body'   : self.body()       ,
+                                     'name'   : self.node.name    ,
+                                     #'returns': self.returns()                 # this is for type hints
+                                     }}
 
     def name(self):
         return self.node.name

diff --git a/osbot_utils/utils/ast/Ast_Load.py b/osbot_utils/utils/ast/Ast_Load.py
@@ -0,0 +1,10 @@
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Load(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Load]'
+
+    def info(self):
+        #self.print()
+        return {'Ast_Load': {}}
diff --git a/osbot_utils/utils/ast/Ast_Module.py b/osbot_utils/utils/ast/Ast_Module.py
@@ -19,16 +19,6 @@ def __repr__(self):
     def source_code(self):
         return ast.unparse(self.module)
 
-    def all_nodes(self):
-        nodes = []
-        for node in ast.walk(self.module):
-            node = self.ast_node(node)
-            nodes.append(node)
-        return nodes
 
     def info(self):                             # todo: see if .info() is the best name for this
-        body_nodes = []
-        for body_node in self.module.body:
-            body_node = self.ast_node(body_node)
-            body_nodes.append(body_node.info())
-        return {'Ast_Module': {'body':body_nodes  } }
+        return {'Ast_Module': {'body':self.body()  } }
diff --git a/osbot_utils/utils/ast/Ast_Name.py b/osbot_utils/utils/ast/Ast_Name.py
@@ -0,0 +1,9 @@
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Name(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Name]'
+
+    def info(self):
+        return {'Ast_Name': {'ctx': self.ast_node(self.node.ctx)}}
diff --git a/osbot_utils/utils/ast/Ast_Node.py b/osbot_utils/utils/ast/Ast_Node.py
@@ -17,20 +17,84 @@ def __init__(self, node):
     def __repr__(self):
         return f"[Ast_Node][????] {self.type}"
 
+    def args(self):
+        return self.ast_nodes(self.node.args)
+
     def ast_node(self, node):
         type_key      = type(node)
         resolved_type = type_registry.resolve(type_key)
         if resolved_type:
            return resolved_type(node)
         return Ast_Node(node)
 
+    def ast_nodes(self, nodes):
+        ast_nodes = []
+        for node in nodes:
+            ast_node = self.ast_node(node)
+            ast_nodes.append(ast_node.info())           # todo: see the use of .info() here (should be better to return the ast_node)
+        return ast_nodes
+
+    def all_nodes(self):
+        nodes = []
+        for node in ast.walk(self.node):
+            node = self.ast_node(node)
+            nodes.append(node)
+        return nodes
+
+    def bases(self):
+        return self.ast_nodes(self.node.bases)
+        # nodes     = self.node.bases
+        # ast_nodes = []
+        # for node in nodes:
+        #     ast_node = self.ast_node(node)
+        #     ast_nodes.append(ast_node.info())
+        # return ast_nodes
+
+    def body(self):
+        return self.ast_nodes(self.node.body)
+
+    def ctx(self):
+        return self.ast_node(self.node.ctx)
+
+    def func(self):
+        return self.ast_node(self.node.func)
+
     def info(self):
         return obj_data(self.node)
 
+    def keys(self):
+        return self.ast_nodes(self.node.keys)
+
+    def keywords(self):
+        return self.ast_nodes(self.node.keywords)
+
     def print(self):
         obj_info(self.node)
         return self
 
     def source_code(self):
         return ast.unparse(self.node)
 
+    def targets(self):
+        return self.ast_nodes(self.node.targets)
+
+    def value(self):
+        return self.ast_node(self.node.value)
+
+    def values(self):
+        return self.ast_nodes(self.node.values)
+
+    def stats(self):
+        types_stats = {}
+        for node in self.all_nodes():
+            type_name = node.type_name
+            if types_stats.get(type_name) is None:
+                types_stats[type_name] = 0
+            types_stats[type_name] += 1
+        stats = {'types': types_stats}
+        return stats
+
+    # def returns(self):                                    # todo: add this when looking at type hints (which is what this is )
+    #     if self.node.returns:
+    #         return self.ast_node(self.node.returns)
+
diff --git a/osbot_utils/utils/ast/Ast_Return.py b/osbot_utils/utils/ast/Ast_Return.py
@@ -6,9 +6,9 @@ def __repr__(self):
         return f'[Ast_Node][Ast_Return]'
 
     def info(self):
-        return {'Ast_Return': {'value': self.value()}}
+        return {'Ast_Return': {'value': self.value().info()}}
 
-    def value(self):
-        return_node = self.ast_node(self.node.value)
-        return return_node.info()
+    # def value(self):
+    #     return_node = self.ast_node(self.node.value)
+    #     return return_node.info()
 
diff --git a/osbot_utils/utils/ast/Ast_Store.py b/osbot_utils/utils/ast/Ast_Store.py
@@ -0,0 +1,10 @@
+from osbot_utils.utils.Dev import pprint
+from osbot_utils.utils.ast.Ast_Node import Ast_Node
+
+class Ast_Store(Ast_Node):
+
+    def __repr__(self):
+        return f'[Ast_Node][Ast_Store]'
+
+    def info(self):
+        return {'Ast_Store': {}}
diff --git a/osbot_utils/utils/ast/__init__.py b/osbot_utils/utils/ast/__init__.py
@@ -5,15 +5,37 @@
 from osbot_utils.helpers.Type_Registry      import type_registry
 from osbot_utils.utils.ast.Ast_Argument     import Ast_Argument
 from osbot_utils.utils.ast.Ast_Arguments    import Ast_Arguments
+from osbot_utils.utils.ast.Ast_Assign import Ast_Assign
+from osbot_utils.utils.ast.Ast_Attribute import Ast_Attribute
+from osbot_utils.utils.ast.Ast_Call import Ast_Call
+from osbot_utils.utils.ast.Ast_Class_Def    import Ast_Class_Def
 from osbot_utils.utils.ast.Ast_Constant     import Ast_Constant
+from osbot_utils.utils.ast.Ast_Dict import Ast_Dict
+from osbot_utils.utils.ast.Ast_Expr import Ast_Expr
 from osbot_utils.utils.ast.Ast_Function_Def import Ast_Function_Def
+from osbot_utils.utils.ast.Ast_Load import Ast_Load
 from osbot_utils.utils.ast.Ast_Module       import Ast_Module
+from osbot_utils.utils.ast.Ast_Name         import Ast_Name
 from osbot_utils.utils.ast.Ast_Return       import Ast_Return
+from osbot_utils.utils.ast.Ast_Store import Ast_Store
 
-type_registry.register(ast.arguments    , Ast_Arguments    )
-type_registry.register(ast.arg          , Ast_Argument     )
-type_registry.register(ast.Constant     , Ast_Constant     )
-type_registry.register(ast.FunctionDef  , Ast_Function_Def )
-type_registry.register(ast.Module       , Ast_Module       )
-type_registry.register(ast.Return       , Ast_Return       )
+ast_types = {
+    ast.Assign      : Ast_Assign       ,
+    ast.Attribute   : Ast_Attribute    ,
+    ast.arg         : Ast_Argument     ,
+    ast.arguments   : Ast_Arguments    ,
+    ast.Call        : Ast_Call         ,
+    ast.ClassDef    : Ast_Class_Def    ,
+    ast.Constant    : Ast_Constant     ,
+    ast.Dict        : Ast_Dict         ,
+    ast.Expr        : Ast_Expr         ,
+    ast.FunctionDef : Ast_Function_Def ,
+    ast.Load        : Ast_Load         ,
+    ast.Module      : Ast_Module       ,
+    ast.Name        : Ast_Name         ,
+    ast.Return      : Ast_Return       ,
+    ast.Store       : Ast_Store
+}
 
+for key, value in ast_types.items():
+    type_registry.register(key, value)
diff --git a/tests/decorators/methods/test_function_type_check.py b/tests/decorators/methods/test_function_type_check.py
@@ -7,12 +7,13 @@ class An_Class:
 
     @function_type_check
     def method_1(self, an_int: int, an_string: str, an_bool : bool = False ) -> int:
-        print('---- in method 1')
+        #print('---- in method 1')
         return 12
 
     @function_type_check
     def method_2(self):
-        print('---- in method 2')
+        #print('---- in method 2')
+        pass
 
     @function_type_check
     def method_3(self,return_me) -> str: