FilesExpand file tree

 master

/

ast_utils.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

130 lines (114 loc) · 3.59 KB

 master

/

ast_utils.py

Top

File metadata and controls

• Code
• Blame

130 lines (114 loc) · 3.59 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

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 typedpython

class CollectNames(ast.NodeVisitor):

_names_ = []

def visit_Name(self, node):

self._names_.append( node )

def collect_names(node):

CollectNames._names_ = names = []

CollectNames().visit( node )

return names

class CollectReturns(ast.NodeVisitor):

_returns_ = []

def visit_Return(self, node):

self._returns_.append( node )

def collect_returns(node):

CollectReturns._returns_ = returns = []

CollectReturns().visit( node )

return returns

def retrieve_vars(body):

local_vars = set()

global_vars = set()

for n in body:

#if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name): ## assignment to local - TODO support `a=b=c`

# local_vars.add( n.targets[0].id )

#elif isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Tuple):

# for c in n.targets[0].elts:

# local_vars.add( c.id )

if isinstance(n, ast.Assign):

user_typedef = None

#targets = list(n.targets)

#targets.reverse()

for i,u in enumerate(n.targets):

if isinstance(u, ast.Name):

if i==0:

if u.id in typedpython.types:

user_typedef = u.id

else:

local_vars.add( u.id )

elif user_typedef:

local_vars.add( '%s=%s' %(user_typedef, u.id) )

user_typedef = None

else:

#add = True ## TODO this should work?!

#for x in local_vars:

# if u.id == x or ('=' in x and x.split('=')[-1] == u.id):

# add = False

# break

#if add:

local_vars.add( u.id )

elif isinstance(u, ast.Tuple):

for uu in u.elts:

if isinstance(uu, ast.Name):

local_vars.add( uu.id )

else:

raise NotImplementedError(uu)

else:

pass ## skips assignment to an attribute `a.x = y`

if user_typedef: ## `int x`

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

local_vars.add( '%s=%s' %(user_typedef, n.value.id))

elif isinstance(n.value, ast.Num):

local_vars.add( '%s=%s' %(user_typedef, n.value.n))

else:

raise SyntaxError(n.value)

elif isinstance(n, ast.Global):

global_vars.update( n.names )

elif hasattr(n, 'body') and not isinstance(n, ast.FunctionDef):

# do a recursive search inside new block except function def

l, g = retrieve_vars(n.body)

local_vars.update(l)

global_vars.update(g)

if hasattr(n, 'orelse'):

l, g = retrieve_vars(n.orelse)

local_vars.update(l)

global_vars.update(g)

return local_vars, global_vars

def retrieve_properties(body):

props = set()

for n in body:

if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Attribute) and isinstance(n.targets[0].value, ast.Name) and n.targets[0].value.id == 'self':

props.add( n.targets[0].attr )

elif hasattr(n, 'body') and not isinstance(n, ast.FunctionDef):

props.update( retrieve_properties(n.body) )

return props

def inspect_function( node ):

local_vars, global_vars = retrieve_vars(node.body)

local_vars = local_vars - global_vars

for arg in node.args.args:

local_vars.add( arg.id )

names = []

returns = []

for n in node.body:

names.extend( collect_names(n) )

returns.extend( collect_returns(n) )

typedefs = {}

for decorator in node.decorator_list:

if isinstance(decorator, ast.Call) and decorator.func.id == 'typedef':

c = decorator

assert len(c.args) == 0 and len(c.keywords)

for kw in c.keywords:

assert isinstance( kw.value, ast.Name)

typedefs[ kw.arg ] = kw.value.id

info = {

'locals':local_vars,

'globals':global_vars,

'name_nodes':names,

'return_nodes':returns,

'typedefs': typedefs

}

return info

def inspect_method( node ):

info = inspect_function( node )

info['properties'] = retrieve_properties( node.body )

return info