extractDef.py

Go to the documentation of this file.

 
 
 
 
import ast
import os
 
 
 
 
class FunctionDefVisitor(ast.NodeVisitor):
    
    def __init__(self):
        self._defNodes_defNodes=[]
 
    
    def functionNodes(self):
        return self._defNodes_defNodes
 
    
    def visit_FunctionDef(self, node):
        self._defNodes_defNodes.append(node)
 
 
 
 
class FromImport(ast.NodeVisitor):
    
    def __init__(self, currentLevel):
        self._importDict_importDict={}
        self._currentLevel_currentLevel=currentLevel
 
    
    @property
    def importDict(self):
        return self._importDict_importDict
 
    
    def visit_ImportFrom(self, node):
        if node.module is not None:
            module=node.module
            if node.level==0:#若是绝对导入，需考虑层级
                tempLst=module.split('.')
                if len(tempLst)==1:
                    module=''
                elif self._currentLevel_currentLevel in tempLst:
                    index=tempLst.index(self._currentLevel_currentLevel)
                    module='.'.join(tempLst[index+1:])
            
            lst=[{'name':name.name,'alias':name.asname} for name in node.names] #可能会import个多个,from A import a,b,c 
            for dic in lst: #lst中每个元素都是字典
                key=module+'.'+dic['name'] #dic['name']可能是*
                key=key.lstrip('.')
                if dic['alias']:
                    self._importDict_importDict[key]=dic['alias']
                else:
                    self._importDict_importDict[key]=dic['name']
 
 
 
 
class Def2format:
    
    def __init__(self):
        self._prefix_prefix=''
        self._relativePath_relativePath='' #记录包的相对路径，例如numpy/core/func.py
 
    
    @property
    def prefix(self):
        return self._prefix_prefix
 
    
    @property
    def relativePath(self):
        return self._relativePath_relativePath
 
    
    def toFormat(self,filePath):
        s=filePath.rsplit(f"site-packages{os.sep}", 1)[-1]
        self._relativePath_relativePath=s
        s=s.replace('\\','.').replace('/','.')
        pos=s.rfind('.')
        s=s[0:pos]
        self._prefix_prefix=s.replace('\\','.').replace('/','.')
 
 
 
 
class AssignVisitor(ast.NodeVisitor):
    
    def __init__(self):
        self._targetCall_targetCall={}
 
    
    def get_target_call(self):
        return self._targetCall_targetCall
 
    
    def visit_Assign(self,node):
        if isinstance(node.value,ast.Call):
            targetName=ast.unparse(node.targets)
            valueExpr=ast.unparse(node.value)
            self._targetCall_targetCall[targetName]=valueExpr