Python Construct

Continuing from enhancement of ast construction based from the code_component, this week was primarily focused on implementing the remaining Python constructs.

In addition to handling child nodes (e.g., Literals, Imports), the definition for each node are now more dynamic. The nodes are connected through compositions another by using compositions to establish relationship where the part_id in composition matches the id of code_component, as in construct_node>

In the previous construct_node the attributes such as value, elts, name, and other is mostly using the name of code_component, either by manually spliting, extract the call function, or just use the value as it is and instatiating it under ast.Constant

construct_node

def construct_node(component):
    label_ = component.name
    if component.type == 'script':
        return ast.Module(name=component.name, body=[], type_ignores=[])
    elif component.type == 'literal':
        return ast.Constant(value=component.name)
    elif component.type == 'list':
        ast_ = ast.List(elts=[], ctx=ast.Load())
        ast_.label = label_
        return ast_
    elif component.type == 'tuple':
        ast_ = ast.Tuple(elts=[], ctx=ast.Load())
        ast_.label = label_
        return ast_
    elif component.type == 'set':
        ast_ = ast.Set(elts=[])
        ast_.label = label_
        return ast_
    elif component.type == 'dict':
        ast_ = ast.Dict(keys=[], values=[])
        ast_.label = label_
        return ast_
    elif component.type == 'key_value':
        key_ = ''.join(component.name.split(':')[0]).strip()
        value_ = ''.join(component.name.split(':')[1]).strip()
        return ([ast.Constant(value=key_), ast.Constant(value=value_)])

    elif component.type == 'name':
        return ast.Name(id=component.name, ctx=ast.Load())
    elif component.type == 'expr':
        ast_ = ast.Expr(value=None)
        ast_.label = label_
        return ast_

    elif component.type == 'uadd':
        ast_ = ast.UnaryOp(op=ast.UAdd(), operand=None)
        ast_.label = label_
        return ast_
    elif component.type == 'usub':
        ast_ = ast.UnaryOp(op=ast.USub(), operand=None)
        ast_.label = label_
        return ast_
    elif component.type == 'not':
        ast_ = ast.UnaryOp(op=ast.Not(), operand=None)
        ast_.label = label_
        return ast_
    elif component.type == 'invert':
        ast_ = ast.UnaryOp(op=ast.Invert(), operand=None)
        ast_.label = label_
        return ast_

    elif component.type == "add":
        ast_ = ast.BinOp(left=None, op=ast.Add(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "sub":
        ast_ = ast.BinOp(left=None, op=ast.Sub(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "div":
        ast_ = ast.BinOp(left=None, op=ast.Div(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "mult":
        ast_ = ast.BinOp(left=None, op=ast.Mult(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "floordiv":
        ast_ = ast.BinOp(left=None, op=ast.FloorDiv(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "mod":
        ast_ = ast.BinOp(left=None, op=ast.Mod(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "pow":
        ast_ = ast.BinOp(left=None, op=ast.Pow(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "lshift":
        ast_ = ast.BinOp(left=None, op=ast.LShift(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "rshift":
        ast_ = ast.BinOp(left=None, op=ast.RShift(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "bitor":
        ast_ = ast.BinOp(left=None, op=ast.BitOr(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "bitxor":
        ast_ = ast.BinOp(left=None, op=ast.BitXor(), right=None)
        ast_.label = label_
        return ast_
    elif component.type == "bitand":
        ast_ = ast.BinOp(left=None, op=ast.BitAnd(), right=None)
        ast_.label = label_
        return ast_

    elif component.type == "and":
        ast_ = ast.BoolOp(op=ast.And(), values=[])
        ast_.label = label_
        return ast_
    elif component.type == "or":
        ast_ = ast.BoolOp(op=ast.Or(), values=[])
        ast_.label = label_
        return ast_
    elif any(t in ['eq', 'noteq', 'lt', 'lte', 'gt', 'gte', 'is', 'isnot', 'in', 'notin'] for t in component.type.split('.')):
        ops_ = []
        ops_map = {
            'eq': ast.Eq(), 'noteq': ast.NotEq(), 'lt': ast.Lt(), 'lte': ast.LtE(),
            'gt': ast.Gt(), 'gte': ast.GtE(), 'is': ast.Is(), 'isnot': ast.IsNot(),
            'in': ast.In(), 'notin': ast.NotIn()
        }
        for t in component.type.split('.'):
            if t in ops_map.keys():
                ops_.append(ops_map[t])
        ast_ = ast.Compare(left=None, ops=ops_, comparators=[])
        ast_.label = label_
        return ast_

    elif component.type == "call":
        ast_ = ast.Call(func=None, args=[], keywords=[])
        ast_.label = label_
        return ast_
    elif component.type == 'ifexp':
        ast_ = ast.IfExp(test=None, body=None, orelse=None)
        ast_.label = label_
        return ast_
    elif component.type == 'attribute':
        ast_ = ast.Attribute(value=''.join(
            component.name.split('.')[1]), attr=None, ctx=ast.Load())
        ast_.label = label_
        return ast_

    elif component.type == "subscript":
        return ast.Subscript(value=None, slice=None, ctx=ast.Load())
    elif component.type == 'index':
        return ast.Index(value=component.name)
    elif component.type == 'slice':
        ast_ = ast.Slice(lower=None, upper=None, step=None)
        ast_.label = label_
        return ast_
    elif component.type == 'extslice':
        ast_ = ast.ExtSlice(dims=[])
        ast_.label = label_
        return ast_

    elif component.type == 'listcomp':
        ast_ = ast.ListComp(elt=None, generators=[])
        ast_.label = label_
        return ast_
    elif component.type == 'setcomp':
        ast_ = ast.SetComp(elt=None, generators=[])
        ast_.label = label_
        return ast_
    elif component.type == 'genexpcomp':
        ast_ = ast.GeneratorExp(elt=None, generators=[])
        ast_.label = label_
        return ast_
    elif component.type == 'dictcomp':
        ast_ = ast.DictComp(key=None, value=None, generators=[])
        ast_.label = label_
        return ast_
    elif component.type == 'comprehension':
        ast_ = ast.comprehension(target=None, iter=None, ifs=[], is_async=0)
        ast_.label = label_
        return ast_

    elif component.type == 'assign':
        ast_ = ast.Assign(targets=[], value=None)
        ast_.label = label_
        return ast_
    elif component.type == 'ann_assign':
        ast_ = ast.AnnAssign(target=None, annotation=None, simple=None)
        ast_.label = label_
        return ast_
    elif component.type == 'ann_target':
        if '.' in component.name:
            values_ = component.name.split('.')
            ast_ = ast.Attribute(value=ast.Name(
                id=values_[0], ctx=ast.Load()), attr=values_[1], ctx=ast.Load())
            ast_.label = label_
            return ast_
        elif '[' in component.name and ']' in component.name:
            value_, slice_ = component.name.split('[')
            slice_ = slice_.rstrip(']')
            if ':' in slice_:
                parts = slice_.split(':')
                lower_ = parts[0].strip() if parts[0].strip() else None
                upper_ = parts[1].strip() if parts[1].strip() else None
                step_ = parts[2].strip() if len(
                    parts) > 2 and parts[2].strip() else None
                ast_ = ast.Subscript(value=ast.Name(id=value_, ctx=ast.Load()),
                                     slice=ast.Slice(lower=lower_, upper=upper_, step=step_), ctx=ast.Store())
                ast_.label = label_
                return ast_
            else:
                ast_ = ast.Subscript(value=ast.Name(id=value_, ctx=ast.Store()),
                                     slice=ast.Index(value=slice_), ctx=ast.Load())
                ast_.label = label_
                return ast_
        else:
            ast_ = ast.Name(id=component.name, ctx=ast.Load())
            ast_.label = label_
            return ast_
    elif component.type == 'annotation':
        ast_ = ast.Constant(value=component.name)
        ast_.label = label_
        return ast_

    elif component.type == 'aug_assign':
        ops_map = {
            '+=': ast.Add(), '-=': ast.Sub(), '*=': ast.Mult(), '/=': ast.Div(),
            '%=': ast.Mod(), '**=': ast.Pow(), '<<=': ast.LShift(), '>>=': ast.RShift(),
            '&=': ast.BitAnd(), '|=': ast.BitOr(), '^=': ast.BitXor()
        }
        op_ = component.name.split()[1].strip()
        ast_ = ast.AugAssign(target=None, op=ops_map[op_], value=None)
        ast_.label = label_
        return ast_
    elif component.type == 'raise':
        ast_ = ast.Raise(exc=ast.Name(id=component.name, ctx=ast.Load()))
        ast_.label = label_
        return ast_
    elif component.type == 'assert':
        test_ = component.name.split('assert ')[1].split(',')[0].strip()
        msg_ = component.name.split('assert ')[1].split(',')[1].strip() if len(
            component.name.split('assert ')[1].split(',')) > 1 else None
        ast_ = ast.Assert(test=ast.Name(id=test_, ctx=ast.Load()),
                          msg=ast.Constant(value=msg_) if msg_ else None)
        ast_.label = label_
        return ast_
    elif component.type == 'delete':
        targets_ = [ast.Name(id=target.strip(), ctx=ast.Del())
                    for target in component.name.split(',')]
        ast_ = ast.Delete(targets=targets_)
        ast_.label = label_
        return ast_
    elif component.type == "pass":
        return ast.Pass()

    elif component.type == 'import':
        name_ = ' '.join(component.name.split()[1:])
        name_ = name_.split(' as ')[0].split()[0]
        asname_ = component.name.split(
            ' as ')[-1].strip() if ' as ' in component.name else None
        return ast.Import(names=[ast.alias(name=name_, asname=asname_)])
    elif component.type == 'import_from':
        module_ = component.name.split()[1]
        level_ = module_.count('.') if module_.startswith('.') else 0
        names_ = ' '.join(component.name.split()[3:])
        names_ = [n_.strip() for n_ in names_.split(',')]
        alias = [ast.alias(name=n.split(' ')[0], asname=n.split(
            ' as ')[-1].strip() if ' as ' in n else None) for n in names_]
        return ast.ImportFrom(module=module_, names=alias, level=level_)

    elif component.type == "if":
        ast_ = ast.If(test=None, body=[], orelse=[])
        ast_.label = label_
        return ast_
    elif component.type == "for":
        ast_ = ast.For(target=None, iter=None, body=[], orelse=[])
        ast_.label = label_
        return ast_
    elif component.type == 'break':
        return ast.Break()
    elif component.type == 'continue':
        return ast.Continue()
    elif component.type == 'while':
        ast_ = ast.While(test=None, body=[], orelse=[])
        ast_.label = label_
        return ast_
    elif component.type == 'try':
        ast_ = ast.Try(body=[], handlers=[], orelse=[], finalbody=[])
        ast_.label = label_
        return ast_
    elif component.type == 'exception':
        ast_ = ast.ExceptHandler(name=component.name, body=[])
        ast_.label = label_
        return ast_
    elif component.type == 'with':
        ast_ = ast.With(items=[], body=[])
        ast_.label = label_
        return ast_
    elif component.type == 'withitem':
        as_ = ''.join(component.name.split(' as ')[
                      1]) if ' as ' in component.name else None
        return ast.withitem(context_expr=None, optional_vars=as_)

    elif component.type == 'function_def':
        args = [body['name'] for body in results_dict
                if body['type'] == 'arguments' and body['whole_id'] == component.id][0]
        label_lines = []
        for body in results_dict:
            if body['whole_id'] == component.id and body['position'] is not None:
                if body['type'] == 'function_def':
                    label_lines.append(f"def {body['name']}({args}):")
                else:
                    label_lines.append(body['name'])
        label_ = '\n'.join(label_lines)
        ast_ = ast.FunctionDef(name=None, args=None,
                               body=[], decorator_list=[], type_params=[])
        ast_.label = label_
        return ast_
    elif component.type == 'identifier':
        return component.name
    elif component.type == 'lambda_def':
        return ast.Lambda(args=None, body=[])
    elif component.type == 'arguments':
        return ast.arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[],
                             kw_defaults=[], kwarg=None, defaults=[])
    elif component.type == 'param':
        return ast.arg(arg=component.name, annotation=None)
    elif component.type == "return":
        ast_ = ast.Return(value=None)
        ast_.label = label_
        return ast_
    elif component.type == 'yield':
        ast_ = ast.Yield(value=None, ctx=ast.Load())
        ast_.label = label_
        return ast_
    elif component.type == 'yield_from':
        ast_ = ast.YieldFrom(value=None)
        ast_.label = label_
        return ast_
    elif component.type == 'global':
        names_ = component.name.replace('global ', '').strip().split(',')
        ast_ = ast.Global(names=names_)
        ast_.label = label_
        return ast_
    elif component.type == 'nonlocal':
        names_ = component.name.replace('nonlocal ', '').strip().split(',')
        ast_ = ast.Nonlocal(names=names_)
        ast_.label = label_
        return ast_

    elif component.type == 'class_def':
        label_lines = []
        for body in results_dict:
            if body['whole_id'] == component.id and body['position'] is not None:
                if body['type'] == 'function_def':
                    label_lines.append(f"def {body['name']}:")
                else:
                    label_lines.append(body['name'])
        label_ = '\n'.join(label_lines)
        ast_ = ast.ClassDef(name=None, bases=[], keywords=[],
                            body=[], decorator_list=[], type_params=[])
        ast_.label = label_
        return ast_

Nodes defined

Script (script): Creates an ast.Module node representing the script/module.
Literal (literal): Creates an ast.Constant node representing a literal value (e.g., numbers, strings).
List (list): Creates an ast.List node representing a list.
Tuple (tuple): Creates an ast.Tuple node representing a tuple.
Set (set): Creates an ast.Set node representing a set.
Dictionary (dict): Creates an ast.Dict node representing a dictionary.
Key-Value Pair (key_value): Creates two ast.Constant nodes representing a key-value pair in a dictionary.
Name (name): Creates an ast.Name node representing a variable or identifier.
Expression (expr): Creates an ast.Expr node representing an expression.
Unary Operations:
- Unary Plus (uadd): Creates an ast.UnaryOp node with ast.UAdd as the operation.
- Unary Minus (usub): Creates an ast.UnaryOp node with ast.USub as the operation.
- Not (not): Creates an ast.UnaryOp node with ast.Not as the operation.
- Invert (invert): Creates an ast.UnaryOp node with ast.Invert as the operation.
Binary Operations:
- Addition (add): Creates an ast.BinOp node with ast.Add as the operation.
- Subtraction (sub): Creates an ast.BinOp node with ast.Sub as the operation.
- Division (div): Creates an ast.BinOp node with ast.Div as the operation.
- Multiplication (mult): Creates an ast.BinOp node with ast.Mult as the operation.
- Floor Division (floordiv): Creates an ast.BinOp node with ast.FloorDiv as the operation.
- Modulus (mod): Creates an ast.BinOp node with ast.Mod as the operation.
- Power (pow): Creates an ast.BinOp node with ast.Pow as the operation.
- Left Shift (lshift): Creates an ast.BinOp node with ast.LShift as the operation.
- Right Shift (rshift): Creates an ast.BinOp node with ast.RShift as the operation.
- Bitwise OR (bitor): Creates an ast.BinOp node with ast.BitOr as the operation.
- Bitwise XOR (bitxor): Creates an ast.BinOp node with ast.BitXor as the operation.
- Bitwise AND (bitand): Creates an ast.BinOp node with ast.BitAnd as the operation.
Boolean Operations:
- And (and): Creates an ast.BoolOp node with ast.And as the operation.
- Or (or): Creates an ast.BoolOp node with ast.Or as the operation.
Comparisons: Creates an ast.Compare node with appropriate comparison operations combinations that seperate with . (ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE, ast.Is, ast.IsNot, ast.In, ast.NotIn).
Function/Method Call (call): Creates an ast.Call node representing a function/method call.
If Else Ternary Operation (ifexp): Creates an ast.IfExp representing a if b else c expression
Attribute (attribute): Creates an ast.Attribute node representing attribute access.
Subscript Operation (subscript): Creates an ast.Subscript node representing a subscript operation.
List Comprehension (listcomp): Creates an ast.ListComp node representing a list comprehension.
Set Comprehension (setcomp): Creates an ast.SetComp node representing a set comprehension.
Generator Expression Comprehension (genexpcomp): Creates an ast.GeneratorExp node representing a generator expression comprehension.
Dictionary Comprehension (dictcomp): Creates an ast.DictComp node representing a dictionary comprehension.
Comprehension (comprehension): Creates an ast.comprehension node representing a comprehension clause.
Assignment (assign): Creates an ast.Assign node representing an assignment statement.
Annotated Assignment (ann_assign): Creates an ast.AnnAssign node representing an annotated assignment statement.
Augmented Assignment (aug_assign): Creates an ast.AugAssign node representing an augmented assignment statement.
Raise (raise): Creates an ast.Raise node representing a raise statement.
Assert (assert): Creates an ast.Assert node representing an assert statement.
Delete (del): Creates an ast.Delete node representing a delete statement.
Pass (pass): Creates an ast.Pass node representing a pass statement.
Import (import): Creates an ast.Import node representing an import statement.
Import From (import_from): Creates an ast.ImportFrom node representing an import from statement.
If Statement (if): Creates an ast.If node representing an if statement.
While Loop (while): Creates an ast.While node representing a while loop statement.
For Loop (for): Creates an ast.For node representing a for loop statement.
Break (break): Creates an ast.Break node representing a break statement.
Continue (continue): Creates an ast.Continue node representing a continue statement.
With Statement (with): Creates an ast.With node representing a with statement.
With Item (withitem): Creates an ast.withitem node representing an item in a with statement.
Function Definition (function_def): Creates an ast.FunctionDef node representing a function definition.
Lambda Function Definition (lambda_def): Creates an astLambda node representing a lambda function definition.
Arguments (arguments): Creates an ast.arguments node representing arguments of a function.
Argument (param): Creates an ast.arg node representing a single argument.
Return (return): Creates an ast.Return node representing a return statement.
Yield (yield): Creates an ast.Yield node representing a yield statement.
Yield From (yieldfrom): Creates an ast.YieldFrom node representing a yield from statement.
Global (global): Creates an ast.Global node representing a global statement.
Nonlocal (nonlocal): Creates an ast.Nonlocal node representing a nonlocal statement.
Class Definition (class_def): Creates an ast.ClassDef node representing a class definition.

Additionally, I added another atribute for the nodes, which is the label. node.label is going to be the label of the node from ast graph visualization. These label is mainly from the name of code_component, appart from function_def & class_def.

Querying the labels for def nodes

The complicated label for the function_def & class_def are queried from the database. The following query extracts the IDs of code components that are either function definitions (function_def) or class definitions (class_def) for a specific trial. The resulting IDs are used as a subquery for the next query.

def_id = session.query(CodeComponent.id).join(
    Composition,
    (CodeComponent.id == Composition.part_id) &
    (CodeComponent.trial_id == Composition.trial_id)
).filter(
    (CodeComponent.trial_id == trial_id) &
    (CodeComponent.type == 'function_def') | (CodeComponent.type == 'class_def')
).subquery()

Next, the details of the code_component are needed. Below query fetches details of code components related to the previously identified function and class definitions for the same trial. The details include the name, type, whole_id, and position fields from the CodeComponent and Composition tables. We join these tables on matching id and trial_id fields, filter the results to exclude components of type syntax, and to include only those whose whole_id is in the list of IDs from our previous query.

results = session.query(
    CodeComponent.name,
    CodeComponent.type,
    Composition.whole_id,
    Composition.position
).join(
    Composition,
    (CodeComponent.id == Composition.part_id) &
    (CodeComponent.trial_id == Composition.trial_id)
).filter(
    (CodeComponent.trial_id == trial_id) &
    (CodeComponent.type != 'syntax') &
    Composition.whole_id.in_(def_id)
).all()

The query results are then converted into a list of dictionaries for easier manipulation and readability.

results_dict = [
    {
        'name': row.name,
        'type': row.type,
        'whole_id': row.whole_id,
        'position': row.position
    }
    for row in results
]

build_ast

As mentioned before, the relationship between code_component & composition is more dynamic. With the Python constructs coverage in noworkflow, the attribute of the nodes are more defined. Each block handles how to assign or append the attribute of parent node to current iteration node based on the composition type.

def build_ast(code_components, compositions):

    node_dict = {}
    for component in code_components:
        node = construct_node(component)
        if node is None:
            continue
        node_dict[component.id] = node

    for composition in compositions:
        whole_node = node_dict.get(composition.whole_id)
        part_node = node_dict.get(composition.part_id)

        if whole_node is None or part_node is None:
            if composition.extra is None:
                continue

        if isinstance(whole_node, ast.Module):
            whole_node.body.append(part_node)

        elif isinstance(whole_node, ast.List) or (isinstance(whole_node, ast.Tuple) and isinstance(part_node, ast.AST)) or isinstance(whole_node, ast.Set):
            whole_node.elts.append(part_node)
            whole_node.content = []
            composition.whole_id == part_node and whole_node.content.append(
                whole_node.name)
        elif isinstance(whole_node, ast.Dict):
            print(part_node[0])
            print(part_node[1])
            whole_node.keys.append(part_node[0])
            whole_node.values.append(part_node[1])

        elif isinstance(whole_node, ast.Expr):
            whole_node.value = part_node

        elif isinstance(whole_node, ast.UnaryOp):
            whole_node.operand = part_node
        elif isinstance(whole_node, ast.BinOp):
            if composition.type == 'left':
                whole_node.left = part_node
            elif composition.type == 'right':
                whole_node.right = part_node
        elif isinstance(whole_node, ast.BoolOp):
            whole_node.values.append(part_node)
        elif isinstance(whole_node, ast.Compare):
            if composition.type == 'left':
                whole_node.left = part_node
            elif composition.type == '*comparators':
                whole_node.comparators.append(part_node)

        elif isinstance(whole_node, ast.Call):
            if composition.type == 'func':
                whole_node.func = part_node
            elif composition.type == '*args':
                whole_node.args.append(part_node)
            elif composition.type == '*keywords':
                whole_node.keywords.append(part_node)
        elif isinstance(whole_node, ast.IfExp):
            if composition.type == 'test':
                whole_node.test = part_node
            elif composition.type == 'body':
                whole_node.body = part_node
            elif composition.type == 'orelse':
                whole_node.orelse = part_node
        elif isinstance(whole_node, ast.Attribute):
            if composition.type == 'value':
                whole_node.value = part_node
            elif composition.extra is not None:
                extra_ = composition.extra[composition.extra.find(
                    "'") + 1:composition.extra.rfind("'")]
                whole_node.attr = extra_

        elif isinstance(whole_node, ast.Subscript):
            if composition.type == 'value':
                whole_node.value = part_node
            elif composition.type == 'slice':
                if isinstance(part_node, ast.AST):
                    whole_node.slice = part_node
                else:
                    whole_node.slice = ast.Constant(value=part_node)
        elif isinstance(whole_node, ast.Slice):
            if composition.type == 'lower':
                whole_node.lower = part_node
            elif composition.type == 'upper':
                whole_node.upper = part_node
            elif composition.type == 'step':
                whole_node.step = part_node
        elif isinstance(whole_node, ast.Tuple):  # ast.ExtSlice
            if isinstance(part_node, ast.AST):
                whole_node.dims.append(part_node)
            else:
                whole_node.dims.append(ast.Constant(value=part_node))

        elif isinstance(whole_node, ast.ListComp) or isinstance(whole_node, ast.SetComp) or isinstance(whole_node, ast.GeneratorExp):
            if composition.type == 'elt':
                whole_node.elt = part_node
            elif composition.type == '*generators':
                whole_node.generators.append(part_node)
        elif isinstance(whole_node, ast.DictComp):
            if composition.type == 'key_value':
                print(part_node[0])
                print(part_node[1])
                whole_node.key = part_node[0]
                whole_node.value = part_node[1]
            elif composition.type == '*generators':
                whole_node.generators.append(part_node)
        elif isinstance(whole_node, ast.comprehension):
            if composition.type == 'target':
                whole_node.target = part_node
            elif composition.type == 'iter':
                whole_node.iter = part_node
            elif composition.type == '*ifs':
                whole_node.ifs.append(part_node)

        elif isinstance(whole_node, ast.Assign):
            if composition.type == '*targets':
                whole_node.targets.append(part_node)
            elif composition.type == 'value':
                whole_node.value = part_node
        elif isinstance(whole_node, ast.AnnAssign):
            if composition.type == 'annotation':
                whole_node.annotation = part_node
            elif composition.type == 'simple':
                simple_ = composition.extra.split('(')[1].split(')')[0]
                whole_node.simple = simple_
            elif composition.type == 'target':
                whole_node.target = part_node
            elif composition.type == 'value':
                whole_node.value = part_node
        elif isinstance(whole_node, ast.AugAssign):
            if composition.type == 'target':
                whole_node.target = part_node
            elif composition.type == 'value':
                whole_node.value = part_node

        elif isinstance(whole_node, ast.If):
            if composition.type == 'test':
                whole_node.test = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)
            elif composition.type == '*orelse':
                whole_node.orelse.append(part_node)
        elif isinstance(whole_node, ast.For):
            if composition.type == 'target':
                whole_node.target = part_node
            elif composition.type == 'iter':
                whole_node.iter = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)
            elif composition.type == '*orelse':
                whole_node.body.append(part_node)
        elif isinstance(whole_node, ast.While):
            if composition.type == 'test':
                whole_node.test = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)
            elif composition.type == '*orelse':
                whole_node.orelse.append(part_node)
        elif isinstance(whole_node, ast.Try):
            if composition.type == '*body':
                whole_node.body.append(part_node)
            elif composition.type == '*handlers':
                whole_node.handlers.append(part_node)
            elif composition.type == '*orelse':
                whole_node.orelse.append(part_node)
            elif composition.type == '*finalbody':
                whole_node.finalbody.append(part_node)
        elif isinstance(whole_node, ast.ExceptHandler):
            if composition.type == '*body':
                whole_node.body.append(part_node)
        elif isinstance(whole_node, ast.With):
            if composition.type == '*items':
                if isinstance(part_node, ast.withitem):
                    whole_node.items.append(part_node)
                else:
                    whole_node.items[composition.position].context_expr = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)

        elif isinstance(whole_node, ast.FunctionDef):
            if composition.type == 'name_node':
                whole_node.name = part_node
            elif composition.type == 'args':
                if isinstance(part_node, ast.arguments):
                    whole_node.args = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)
            elif composition.type == '*decorator_list':
                whole_node.decorator_list.append(part_node)
        elif isinstance(whole_node, ast.Lambda):
            if composition.type == 'args':
                whole_node.args = part_node
            elif composition.type == 'body':
                whole_node.body = part_node
        elif isinstance(whole_node, ast.Return):
            whole_node.value = part_node
        elif isinstance(whole_node, ast.arguments):
            if composition.type == '*args':
                whole_node.args.append(part_node)
            elif composition.type == 'vararg':
                whole_node.vararg = part_node
            elif composition.type == '*defaults':
                whole_node.defaults.append(part_node)
            elif composition.type == 'kwarg':
                whole_node.kwarg = part_node
            elif composition.type == '*kwonlyargs':
                whole_node.kwonlyargs.append(part_node)
            elif composition.type == '*kw_defaults':
                whole_node.kw_defaults.append(part_node)
            elif composition.type == '*posonlyargs':
                whole_node.posonlyargs.append(part_node)
        elif isinstance(whole_node, ast.Yield) or isinstance(whole_node, ast.YieldFrom):
            whole_node.value = part_node
        elif isinstance(whole_node, ast.ClassDef):
            if composition.type == 'name_node':
                whole_node.name = part_node
            elif composition.type == '*body':
                whole_node.body.append(part_node)

    return node_dict

ast_to_dot

The label for the AST is displayed as mention before. The sample of the output label on Name node:

Name
x

def ast_to_dot(node):
    """Converts AST node to DOT format for Graphviz."""
    def node_label(node):
        label = type(node).__name__
        if isinstance(node, ast.Module):
            return f"{label}\n{node.name}"
        elif isinstance(node, ast.Constant):
            return f"{label}\n{node.value}"
        elif isinstance(node, ast.Name):
            return f"{label}\n{node.id}"
        elif isinstance(node, ast.Import):
            return f"Import\nimport {node.names[0].name}"
        elif isinstance(node, ast.ImportFrom):
            names_ = ', '.join([names.name for names in node.names])
            return f"ImportFrom\nfrom {node.module} import {names_}"
        elif isinstance(node, ast.alias):
            return f"alias\n{node.name} as {node.asname}" if node.asname else f"alias\n{node.name}"
        elif isinstance(node, ast.FunctionDef):
            args = ', '.join(arg.arg for arg in node.args.args)
            return f"{label}\ndef {node.name}({args}):\n{node.label}"
        elif isinstance(node, ast.arg):
            return f"{label}\n{node.arg}"
        elif isinstance(node, ast.ClassDef):
            return f"{label}\nclass {node.name}:\n{node.label}"
        elif hasattr(node, 'label') and isinstance(node, ast.AST):
            return f"{label}\n{node.label}"
        else:
          return f"{label}"

    def node_id(node):
        return f"node{str(id(node))}"

    dot = graphviz.Digraph()
    dot.attr(rankdir='TB', nodesep='0.75', ranksep='0.75')
    root_id = node_id(node)

    def visit(node, parent=None):
        if not isinstance(node, ast.AST) or isinstance(node, ast.Load) or isinstance(node, ast.Store):
            return

        node_name = node_id(node)
        dot.node(node_name, node_label(node), margin='0.1,0.1', width='0.2', height='0.2', fontsize='10')

        if parent:
            dot.edge(parent, node_name, minlen='1', arrowsize='0.5')

        for child_name, child_node in ast.iter_fields(node):
            if isinstance(child_node, list):
                for child in child_node:
                    if isinstance(child, ast.AST):
                        visit(child, node_name)
            elif isinstance(child_node, ast.AST):
                visit(child_node, node_name)

    visit(node)
    return dot

The labels are no longer display the AST Object but instead use the node.label. This is more apparent for nodes that has multiple sub-nodes, such as if, for, function_def, etc.

Correct AST Graph

Nodes Issue

While most of the python construct pre 3.12 has been defined properly, some constructs have some incompability with the noworkflow. For more details on the nodes, refer to below issues.

With the ast for the code_components is finished, now next task is to integrate the now vis

Tags: GSoC Weekly Report

GSoC 2024 - Week 4 Progress

code_component AST graph fully implemented