vim-python-docstring/python/asthelper.py at master · pixelneo/vim-python-docstring · GitHub

History

130 lines (103 loc) · 3.82 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

import ast

import sys

from itertools import chain

class RaiseNameCollector(ast.NodeVisitor):

def __init__(self):

self.data = set()

super().__init__()

def visit_Call(self, node):

self.data.add(node.func.id)

class AttributeCollector(ast.NodeVisitor):

def __init__(self, instance_name):

self.instance_name = instance_name

self.data = {}

super().__init__()

def visit_Attribute(self, node):

if isinstance(node.value, ast.Name):

if node.value.id == self.instance_name:

self.data[node.attr] = None

else:

self.generic_visit(node)

class ClassInstanceNameExtractor(ast.NodeVisitor):

def __init__(self):

self.instance_name = "self" # default

self.set = False

super().__init__()

def visit_FunctionDef(self, node):

if node.name == "__init__":

self.instance_name = node.args.args[0].arg

self.set = True

elif not self.set:

self.instance_name = node.args.args[0].arg

def generic_visit(self, node):

if not self.set:

super().generic_visit(node)

class ClassVisitor(ast.NodeVisitor):

def __init__(self, instance_name):

super().__init__()

self.attributes = {}

self.instance_name = instance_name

def visit_Assign(self, node):

ac = AttributeCollector(self.instance_name)

for target in node.targets:

ac.visit(target)

self.attributes |= ac.data

def visit_AnnAssign(self, node):

ac = AttributeCollector(self.instance_name)

ac.visit(node.target)

self.attributes |= ac.data

class MethodVisitor(ast.NodeVisitor):

"""Gathers information about a method

Attributes:

arguments: arguments of the method

parent: indicated whether this method is inside another

raises: set of raised exceptions

returns: True if method returns

yields: True is method yields

"""

def __init__(self, parent=True):

self.parent = parent

self.arguments = []

self.raises = set()

self.returns = False

self.yields = False

super().__init__()

def _handle_functions(self, node):

new_visitor = MethodVisitor(parent=False)

new_visitor.generic_visit(node)

self.raises |= new_visitor.raises

if self.parent:

for arg in chain(node.args.args, node.args.kwonlyargs):

type_hint = None

if arg.annotation is not None:

# ast.unparse doesn't work for python <= 3.8

if sys.version_info[0] == 3 and sys.version_info[1] <= 8:

from unparse import Unparser

from io import StringIO

v = StringIO()

Unparser(arg.annotation, file=v)

type_hint = v.getvalue()

else:

type_hint = ast.unparse(arg.annotation)

self.arguments.append({"arg": arg.arg, "type": type_hint})

if len(self.arguments) > 0 and (

self.arguments[0]["arg"] == "self" or self.arguments[0]["arg"] == "cls"

self.arguments.pop(0)

self.returns = new_visitor.returns

self.yields = new_visitor.yields

def visit_Raise(self, node):

r = RaiseNameCollector()

r.visit(node)

self.raises |= r.data

super().generic_visit(node)

def visit_Yield(self, node):

self.yields = True

super().generic_visit(node)

def visit_Return(self, node):

self.returns = True

super().generic_visit(node)

def visit_FunctionDef(self, node):

self._handle_functions(node)

def visit_AsyncFunctionDef(self, node):

self._handle_functions(node)

vim-python-docstring/python/asthelper.py at master · pixelneo/vim-python-docstring · GitHub

Expand file tree

Breadcrumbs

asthelper.py

Latest commit

History

Breadcrumbs

asthelper.py

File metadata and controls