utils.ModuleParser

  1import os
  2import astroid
  3from astroid import NodeNG, FunctionDef, ClassDef
  4
  5from data.SourceFile import SourceFile
  6from data.SourceClass import SourceClass
  7from data.SourceFunction import SourceFunction
  8from data.SourceType import SourceType
  9from data.SourceVariable import SourceVariable
 10
 11
 12def annotation_to_type(node: NodeNG | None) -> SourceType:
 13    """
 14    Generates a type string from annotation nodes.
 15
 16    :param node: The root annotation node.
 17    :return: A human-readable type string.
 18    """
 19    string_annotation: str = ''
 20    type_set: set[str] = set()
 21    if isinstance(node, astroid.Name):
 22        string_annotation = node.name
 23        type_set.add(node.name)
 24    elif isinstance(node, astroid.Subscript) and isinstance(node.value, astroid.Name):
 25        if isinstance(node.slice, astroid.Tuple):
 26            string_annotation = node.value.name + '['
 27            for child_node in node.slice.elts:
 28                inner_annotation = annotation_to_type(child_node)
 29                string_annotation += str(inner_annotation) + ', '
 30                type_set.update(inner_annotation.dependencies)
 31            string_annotation = string_annotation[:-2]
 32            string_annotation += ']'
 33        elif isinstance(node.slice, (astroid.Attribute, astroid.Name)):
 34            inner_annotation = annotation_to_type(node.slice)
 35            string_annotation = node.value.name + '[' + str(inner_annotation) + ']'
 36            type_set.update(inner_annotation.dependencies)
 37    # pylint:disable-next=confusing-consecutive-elif
 38    elif isinstance(node, astroid.Attribute):
 39        string_annotation = node.attrname
 40        type_set.add(node.attrname)
 41    elif isinstance(node, astroid.Const):
 42        string_annotation = str(node.value)
 43        type_set.add(str(node.value))
 44    elif isinstance(node, astroid.BinOp):
 45        left = annotation_to_type(node.left)
 46        right = annotation_to_type(node.right)
 47        string_annotation = str(left) + ' | ' + str(right)
 48        type_set.update(left.dependencies)
 49        type_set.update(right.dependencies)
 50    return SourceType(string_annotation, type_set)
 51
 52
 53def get_function(node: FunctionDef) -> SourceFunction:
 54    """
 55    Generates a source function based on the input node.
 56
 57    :param node: The input node.
 58    :return: A source function.
 59    """
 60    annotation = annotation_to_type(node.returns)
 61    function: SourceFunction = SourceFunction()
 62    function.name = node.name
 63    function.returns = annotation
 64    if node.decorators is not None:
 65        for decorator in node.decorators.nodes:
 66            if isinstance(decorator, astroid.Name) and \
 67                    decorator.name == 'staticmethod':
 68                function.static = True
 69    for i, arg in enumerate(node.args.args):
 70        annotation = annotation_to_type(node.args.annotations[i])
 71        try:
 72            default = str(node.args.default_value(arg.name).value)
 73        except astroid.exceptions.NoDefault:
 74            default = None
 75        except AttributeError:
 76            default = None
 77        function.params.append(
 78            SourceVariable(
 79                arg.name,
 80                annotation,
 81                False,
 82                default
 83            )
 84        )
 85    return function
 86
 87
 88def get_class(node: ClassDef) -> SourceClass:
 89    """
 90    Generates a source class based on the input node.
 91
 92    :param node: The input node.
 93    :return: A source class.
 94    """
 95    source_class: SourceClass = SourceClass()
 96    source_class.name = node.name
 97    source_class.bases = [base.name for base in node.bases if isinstance(base, astroid.Name)]
 98    for child_node in node.body:
 99        if isinstance(child_node, FunctionDef):
100            source_class.methods.append(get_function(child_node))
101            for child_child_node in child_node.body:
102                if isinstance(child_child_node, astroid.AnnAssign) and \
103                        isinstance(child_child_node.target, astroid.AssignAttr):
104                    source_class.variables.append(
105                        SourceVariable(
106                            child_child_node.target.attrname,
107                            annotation_to_type(child_child_node.annotation)
108                        )
109                    )
110        elif isinstance(child_node, astroid.AnnAssign) and \
111                isinstance(child_node.target, astroid.AssignName):
112            source_class.variables.append(
113                SourceVariable(
114                    child_node.target.name,
115                    annotation_to_type(child_node.annotation),
116                    True
117                )
118            )
119    return source_class
120
121
122# pylint:disable-next=too-complex
123def get_module_info(
124        root: str,
125        file_path: str,
126        known_modules: list[str],
127        with_external_dependencies: bool = False
128) -> SourceFile:
129    """
130    Generates a source file based on the provided file path.
131
132    :param root: The root path of the analysis.
133    :param file_path: The file path for which to generate the source file.
134    :param known_modules: All known local modules.
135    :param with_external_dependencies: Whether external dependencies should be included.
136    :return: A matching source file.
137    """
138    module = astroid.MANAGER.ast_from_file(file_path)
139    path_modules = file_path[len(root):].split('.')[-2].split(os.sep)
140    source_file = SourceFile()
141    source_file.name = '.'.join(path_modules)
142
143    for node in module.body:
144        if isinstance(node, astroid.ClassDef):
145            source_file.classes.append(get_class(node))
146        elif isinstance(node, FunctionDef):
147            source_file.functions.append(get_function(node))
148        elif isinstance(node, astroid.AnnAssign) and \
149                isinstance(node.target, astroid.AssignName):
150            source_file.variables.append(
151                SourceVariable(node.target.name, annotation_to_type(node.annotation))
152            )
153        elif isinstance(node, astroid.Import):
154            if with_external_dependencies:
155                for node_name in node.names:
156                    source_file.imports.append(node_name[0])
157        # pylint:disable-next=confusing-consecutive-elif
158        elif isinstance(node, astroid.ImportFrom):
159            if node.level == 1 or node.modname in known_modules or with_external_dependencies:
160                for node_name in node.names:
161                    source_file.imports.append(
162                        ('.'.join(path_modules[:-1]) + '.' if node.level == 1 else '') +
163                        node.modname + '.' + node_name[0]
164                    )
165    return source_file