cyyeh/py-code-analyzer

1"""CodeImportsAnalyzer uses the ast module from Python's standard library2to get what modules are imported in given python files, then uses networkx to generate imports graph3"""4import ast5import asyncio6 7import aiohttp8import pybase649 10from .graph_analyzer import GraphAnalyzer11 12 13def construct_fetch_program_text_api_url(api_url):14    import os15 16    # to increase api rate limiting17    # https://docs.github.com/en/rest/overview/resources-in-the-rest-api#rate-limiting18    USER = os.environ.get("USER", "")19    PERSONAL_ACCESS_TOKEN = os.environ.get("PERSONAL_ACCESS_TOKEN", "")20 21    if USER and PERSONAL_ACCESS_TOKEN:22        protocol, api_url_components = api_url.split("://")23        new_api_url_components = f"{USER}:{PERSONAL_ACCESS_TOKEN}@{api_url_components}"24        return f"{protocol}://{new_api_url_components}"25    else:26        return api_url27 28 29async def get_program_text(session, python_file):30    # about Retry-After31    # https://docs.github.com/en/rest/guides/best-practices-for-integrators#dealing-with-secondary-rate-limits32    async with session.get(33        construct_fetch_program_text_api_url(python_file["url"]),34        headers={"Accept": "application/vnd.github.v3+json", "Retry-After": "5"},35    ) as response:36        if response.status == 200:37            data = await response.json()38            if data["encoding"] == "base64":39                return data["content"], python_file["path"]40            else:41                print(42                    f"WARNING: {python_file['path']}'s encoding is {data['encoding']}, not base64"43                )44 45 46class CodeImportsAnalyzer:47    class _NodeVisitor(ast.NodeVisitor):48        def __init__(self, imports):49            self.imports = imports50 51        def visit_Import(self, node):52            for alias in node.names:53                self.imports[-1]["imports"].append(54                    {"module": None, "name": alias.name, "level": -1}55                )56            self.generic_visit(node)57 58        def visit_ImportFrom(self, node):59            for alias in node.names:60                self.imports[-1]["imports"].append(61                    {"module": node.module, "name": alias.name, "level": node.level}62                )63            self.generic_visit(node)64 65    def __init__(self, python_files):66        self.python_imports = []67        self.graph_analyzer = GraphAnalyzer(is_directed=True)68        self.python_files = python_files69        self._node_visitor = CodeImportsAnalyzer._NodeVisitor(self.python_imports)70 71    async def parse_python_files(self):72        async with aiohttp.ClientSession() as session:73            tasks = []74            for python_file in self.python_files:75                tasks.append(76                    asyncio.ensure_future(get_program_text(session, python_file))77                )78 79            results = await asyncio.gather(*tasks)80            if results:81                for base64_program_text, python_file_path in results:82                    if base64_program_text:83                        self.python_imports += [84                            {85                                "file_name": python_file_path.split("/")[-1],86                                "file_path": python_file_path,87                                "imports": [],88                            }89                        ]90                        program = pybase64.b64decode(base64_program_text)91                        tree = ast.parse(program)92                        self._node_visitor.visit(tree)93 94    def generate_imports_graph(self):95        # TODO: thought on how to improve the graph generation logic96        # generate a dictionary of lists data structure97        # generate a graph based on a dictionary of lists98 99        for python_import in self.python_imports:100            _nodes = python_import["file_path"].split("/")101            if len(_nodes):102                # generate graph based on file_path103                # node/edge relationship means file/folder structure104                if len(_nodes) > 1:105                    # make last node and second last node as one node106                    # to solve the issue of duplicated file names using only last node107                    if len(_nodes) >= 3:108                        _nodes[-2] = _nodes[-2] + "/" + _nodes[-1]109                        del _nodes[-1]110                    self.graph_analyzer.add_edges_from_nodes(_nodes)111                else:112                    self.graph_analyzer.add_node(_nodes[0])113 114                # generate graph based on imported modules in each file115                if python_import["file_name"] != "__init__.py":116                    for _import in python_import["imports"]:117                        if _import["module"] is None:118                            _import_names = _import["name"].split(".")119                            _new_nodes = _import_names + [_nodes[-1]]120                            self.graph_analyzer.add_edges_from_nodes(_new_nodes)121                        else:122                            _import_names = _import["module"].split(".") + [123                                _import["name"]124                            ]125                            _new_nodes = _import_names + [_nodes[-1]]126                            self.graph_analyzer.add_edges_from_nodes(_new_nodes)127 128        return self.graph_analyzer.graph129 130    def report(self):131        from pprint import pprint132 133        pprint(self.python_imports)134