Enhanced PythonEvaluator to accept concepts

core/ast/nodes.py

@@ -0,0 +1,110 @@
+from core.builtin_concepts import BuiltinConcepts, ListConcept
+from core.concept import Concept
+import ast
+
+import logging
+
+log = logging.getLogger(__name__)
+
+
+class NodeParent:
+    """
+    Class that represent the ancestor of a Node
+    For example, the 'For' nodes has three fields (target, iter and body)
+    So, for a node under For.iter
+    node -> For
+    field -> iter
+    """
+
+    def __init__(self, node, field):
+        self.node = node
+        self.field = field
+
+    def __repr__(self):
+        if self.node is None:
+            return None
+
+        if self.field is None:
+            return self.node.get_node_type()
+
+        return self.node.get_node_type() + "." + self.field
+
+    def __eq__(self, other):
+        # I can compare with type for simplification
+        if isinstance(other, tuple):
+            return self.node.get_node_type() == other[0] and self.field == other[1]
+
+        # normal equals implementation
+        if not isinstance(other, NodeParent):
+            return False
+
+        return self.node.get_node_type() == other.node.get_node_type() and self.field == other.field
+
+    def __hash__(self):
+        return hash((self.node.get_node_type(), self.field))
+
+
+class NodeConcept(Concept):
+    def __init__(self, key, parent: NodeParent):
+        super().__init__(key, True, False, key)
+        self.parent = parent
+
+    def get_node_type(self):
+        return self.key
+
+
+class GenericNodeConcept(NodeConcept):
+    def __init__(self, node_type, parent):
+        super().__init__(BuiltinConcepts.GENERIC_NODE, parent)
+        self.node_type = node_type
+
+    def __repr__(self):
+        return "Generic:" + self.node_type
+
+    def get_node_type(self):
+        return self.node_type
+
+    def get_value(self):
+        if self.node_type == "Name":
+            return self.get_prop("id")
+
+        if self.node_type == "arg":
+            return self.get_prop("arg")
+
+        return self.body
+
+
+class IdentifierConcept(NodeConcept):
+    def __init__(self, parent, name):
+        super().__init__(BuiltinConcepts.IDENTIFIER_NODE, parent)
+        self.body = name
+
+
+def transform(node):
+    """
+    Transform Python AST node into concept nodes
+    for better usage
+    :param node:
+    :return:
+    """
+
+    def _transform(node, parent):
+        node_type = node.__class__.__name__
+        concept = GenericNodeConcept(node_type, parent).init_key()
+        for field in node._fields:
+            if not hasattr(node, field):
+                continue
+
+            value = getattr(node, field)
+            if isinstance(value, list):
+                lst = ListConcept().init_key()
+                for i in value:
+                    lst.append(_transform(i, NodeParent(concept, field)))
+                concept.set_prop(field, lst)
+            elif isinstance(value, ast.AST):
+                concept.set_prop(field, _transform(value, NodeParent(concept, field)))
+            else:
+                concept.set_prop(field, value)
+        return concept
+
+    return _transform(node, None)