'Extracting the nodes data in AST after parsing the code

import ast

code = """
import sklearn as sk
import numpy as np
import pandas as pd
import matplolib.pyplot as plt
"""

element = ast.parse(code, "<string>", mode='exec')
print(ast.dump(element))

print(element.body)

I am getting output as follows:

Module(body=[Import(names=[alias(name='sklearn', asname='sk')]), Import(names=[alias(name='numpy', asname='np')]), Import(names=[alias(name='pandas', asname='pd')]), Import(names=[alias(name='matplolib.pyplot', asname='plt')])], type_ignores=[])

[<_ast.Import object at 0x7fc18c135910>, <_ast.Import object at 0x7fc18c135bb0>, <_ast.Import object at 0x7fc18c1355b0>, <_ast.Import object at 0x7fc18c0d9550>]

I would like to go through the nodes and extract the name and asname information in Names to store them in imports. I am trying this code whereas it is not working.

modules = {}
imports = []

if isinstance(element, ast.Import) or isinstance(element, ast.ImportFrom):
    element: ast.Import
    if hasattr(element, "module"):
        if element.module not in modules:
            modules = {'imports': []}
        for name in element.names:
            if names.asname != None:
                imports.append(name.asname)
                modules[element.module]["imports"].append(
                    {"name": names.name, "alias": names.asname})
            else:
                imports.append(names.name)
                modules[element.module]["imports"].append(names.name)

print(modules, imports)

I am getting output as, {} [] empty data structures



Solution 1:[1]

You can traverse the ast nodes and extract the module names and aliases from any existing ast.Import objects:

import ast
code = """
import sklearn as sk
import numpy as np
import pandas as pd
import matplolib.pyplot as plt
"""
r = [{'name':i.name, 'alias':i.asname} for j in ast.walk(ast.parse(code))
       for i in getattr(j, 'names', []) if isinstance(j, ast.Import)]

Output:

[{'name': 'sklearn', 'alias': 'sk'}, {'name': 'numpy', 'alias': 'np'}, {'name': 'pandas', 'alias': 'pd'}, {'name': 'matplolib.pyplot', 'alias': 'plt'}]

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 Ajax1234