├── .gitignore ├── .travis.yml ├── README.rst ├── astmonkey ├── __init__.py ├── tests │ ├── __init__.py │ ├── test_transformers.py │ ├── test_utils.py │ └── test_visitors.py ├── transformers.py ├── utils.py └── visitors.py ├── examples ├── edge-graph-node-visitor.py ├── edge-graph.png ├── graph.png ├── graph_node_visitor.py ├── is_docstring.py ├── parent_node_transformer.py └── source_generator_node_visitor.py ├── setup.py └── tox.ini /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.egg* 3 | .*.swp 4 | .*.swo 5 | .DS_Store 6 | .tox 7 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | 3 | matrix: 4 | include: 5 | - python: 2.7 6 | - python: 3.4 7 | - python: 3.5 8 | - python: 3.6 9 | - python: 3.7 10 | - python: 3.8 11 | install: 12 | # python 2.6 and 3.3 compatibility 13 | - if [[ $TRAVIS_PYTHON_VERSION == 2.6 ]]; then pip install pycparser==2.14 cryptography==2.1.4; fi 14 | - if [[ $TRAVIS_PYTHON_VERSION == 3.3 ]]; then pip install virtualenv==15.1.0 setuptools==39.2.0; fi 15 | - pip install tox-travis 16 | - pip install coveralls 17 | script: 18 | - tox 19 | after_success: 20 | coveralls 21 | deploy: 22 | provider: pypi 23 | user: khalas 24 | password: 25 | secure: "rrvTrcLSuZvIeMxP2A84WYB4I1Aea6ymGKuFInSjxReKhG+S9hmSuMyzZASbpGQftZUScCPI4SP3jJQMnEnUJC2jDp+lAjakxDS/E6NMtsIq2vboFG7sejo5LLFCk0vvFjlzZWagYohbMDdqiRT2JEqVnkfk5t2dinA2mgCucDg=" 26 | on: 27 | tags: true 28 | branch: master 29 | python: '3.6' 30 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ========= 2 | astmonkey 3 | ========= 4 | 5 | |Python Versions| |Build Status| |Coverage Status| |Code Climate| 6 | 7 | ``astmonkey`` is a set of tools to play with Python AST. 8 | 9 | Installation 10 | ------------ 11 | 12 | You can install ``astmonkey`` from PyPI: 13 | 14 | :: 15 | 16 | $ pip install astmonkey 17 | 18 | If you want to have latest changes you should clone this repository and use ``setup.py``: 19 | 20 | :: 21 | 22 | $ git clone https://github.com/mutpy/astmonkey.git 23 | $ cd astmonkey 24 | $ python setup.py install 25 | 26 | visitors.SourceGeneratorNodeVisitor 27 | ----------------------------------- 28 | 29 | This visitor allow AST to Python code generation. It was originally written by 30 | Armin Ronacher (2008, license BSD) as ``codegen.py`` module. ``astmonkey`` version 31 | fixes few bugs and it has good code coverage. 32 | 33 | Example usage: 34 | 35 | :: 36 | 37 | import ast 38 | from astmonkey import visitors 39 | 40 | code = 'x = (y + 1)' 41 | node = ast.parse(code) 42 | generated_code = visitors.to_source(node) 43 | 44 | assert(code == generated_code) 45 | 46 | transformers.ParentChildNodeTransformer 47 | --------------------------------------- 48 | 49 | This transformer adds few fields to every node in AST: 50 | 51 | * ``parent`` - link to parent node, 52 | * ``parents`` - list of all parents (only ``ast.expr_context`` nodes have more than one parent node, in other causes this is one-element list), 53 | * ``parent_field`` - name of field in parent node including child node, 54 | * ``parent_field_index`` - parent node field index, if it is a list. 55 | * ``children`` - link to children nodes. 56 | 57 | Example usage: 58 | 59 | :: 60 | 61 | import ast 62 | from astmonkey import transformers 63 | 64 | node = ast.parse('x = 1') 65 | node = transformers.ParentChildNodeTransformer().visit(node) 66 | 67 | assert(node == node.body[0].parent) 68 | assert(node.body[0].parent_field == 'body') 69 | assert(node.body[0].parent_field_index == 0) 70 | assert(node.body[0] in node.children) 71 | 72 | visitors.GraphNodeVisitor 73 | ------------------------- 74 | 75 | This visitor creates Graphviz graph from Python AST (via ``pydot``). Before you use 76 | ``GraphNodeVisitor`` you need to add parents links to tree nodes (with 77 | ``ParentChildNodeTransformer``). 78 | 79 | Example usage: 80 | 81 | :: 82 | 83 | import ast 84 | from astmonkey import visitors, transformers 85 | 86 | node = ast.parse('def foo(x):\n\treturn x + 1') 87 | node = transformers.ParentChildNodeTransformer().visit(node) 88 | visitor = visitors.GraphNodeVisitor() 89 | visitor.visit(node) 90 | 91 | visitor.graph.write_png('graph.png') 92 | 93 | Produced ``graph.png`` (you need to have installed ``graphviz`` binaries if you want generate 94 | images): 95 | 96 | .. image:: examples/graph.png 97 | 98 | utils.is_docstring 99 | ------------------ 100 | 101 | This routine checks if target node is a docstring. Before you use 102 | ``is_docstring`` you need to add parents links to tree nodes (with 103 | ``ParentChildNodeTransformer``). 104 | 105 | Example usage: 106 | 107 | :: 108 | 109 | import ast 110 | from astmonkey import utils, transformers 111 | 112 | node = ast.parse('def foo(x):\n\t"""doc"""') 113 | node = transformers.ParentChildNodeTransformer().visit(node) 114 | 115 | docstring_node = node.body[0].body[0].value 116 | assert(not utils.is_docstring(node)) 117 | assert(utils.is_docstring(docstring_node)) 118 | 119 | 120 | License 121 | ------- 122 | 123 | Copyright [2013] [Konrad Hałas] 124 | 125 | Licensed under the Apache License, Version 2.0 (the "License"); 126 | you may not use this file except in compliance with the License. 127 | You may obtain a copy of the License at 128 | 129 | http://www.apache.org/licenses/LICENSE-2.0 130 | 131 | Unless required by applicable law or agreed to in writing, software 132 | distributed under the License is distributed on an "AS IS" BASIS, 133 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 134 | See the License for the specific language governing permissions and 135 | limitations under the License. 136 | 137 | .. |Python Versions| image:: https://img.shields.io/pypi/pyversions/astmonkey.svg 138 | :target: https://github.com/mutpy/astmonkey 139 | .. |Build Status| image:: https://travis-ci.org/mutpy/astmonkey.png 140 | :target: https://travis-ci.org/mutpy/astmonkey 141 | .. |Coverage Status| image:: https://coveralls.io/repos/github/mutpy/astmonkey/badge.svg?branch=master 142 | :target: https://coveralls.io/github/mutpy/astmonkey?branch=master 143 | .. |Code Climate| image:: https://codeclimate.com/github/mutpy/astmonkey/badges/gpa.svg 144 | :target: https://codeclimate.com/github/mutpy/astmonkey 145 | -------------------------------------------------------------------------------- /astmonkey/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.3.6' 2 | -------------------------------------------------------------------------------- /astmonkey/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mutpy/astmonkey/16b1817fd022b60b1369281af74b465db6db85ff/astmonkey/tests/__init__.py -------------------------------------------------------------------------------- /astmonkey/tests/test_transformers.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | try: 4 | import unittest2 as unittest 5 | except ImportError: 6 | import unittest 7 | import ast 8 | 9 | from astmonkey import transformers 10 | 11 | 12 | class TestParentChildNodeTransformer(object): 13 | @pytest.fixture 14 | def transformer(self): 15 | return transformers.ParentChildNodeTransformer() 16 | 17 | def test_module_node(self, transformer): 18 | node = ast.parse('') 19 | 20 | transformed_node = transformer.visit(node) 21 | 22 | assert transformed_node.parent is None 23 | assert transformed_node.children == [] 24 | 25 | def test_non_module_node(self, transformer): 26 | node = ast.parse('x = 1') 27 | 28 | transformed_node = transformer.visit(node) 29 | 30 | assign_node = transformed_node.body[0] 31 | assert transformed_node == assign_node.parent 32 | assert assign_node.parent_field == 'body' 33 | assert assign_node.parent_field_index == 0 34 | assert transformed_node.children == [assign_node] 35 | assert len(assign_node.children) == 2 36 | 37 | def test_expr_context_nodes(self, transformer): 38 | node = ast.parse('x = 1\nx = 2') 39 | 40 | transformer.visit(node) 41 | 42 | ctx_node = node.body[0].targets[0].ctx 43 | first_name_node = node.body[0].targets[0] 44 | second_name_node = node.body[1].targets[0] 45 | assert first_name_node in ctx_node.parents 46 | assert second_name_node in ctx_node.parents 47 | assert ctx_node in first_name_node.children 48 | assert ctx_node in second_name_node.children 49 | -------------------------------------------------------------------------------- /astmonkey/tests/test_utils.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import unittest 3 | 4 | from astmonkey import utils, transformers 5 | 6 | 7 | class TestIsDocstring(unittest.TestCase): 8 | def test_non_docstring_node(self): 9 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('')) 10 | 11 | assert not utils.is_docstring(node) 12 | 13 | def test_module_docstring_node(self): 14 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('"""doc"""')) 15 | 16 | assert utils.is_docstring(node.body[0].value) 17 | 18 | def test_function_docstring_node(self): 19 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('def foo():\n\t"""doc"""')) 20 | 21 | assert utils.is_docstring(node.body[0].body[0].value) 22 | 23 | def test_class_docstring_node(self): 24 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('class X:\n\t"""doc"""')) 25 | 26 | assert utils.is_docstring(node.body[0].body[0].value) 27 | -------------------------------------------------------------------------------- /astmonkey/tests/test_visitors.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import sys 3 | 4 | import pytest 5 | 6 | try: 7 | import unittest2 as unittest 8 | except ImportError: 9 | import unittest 10 | import ast 11 | from astmonkey import visitors, transformers, utils 12 | 13 | 14 | class TestGraphNodeVisitor(object): 15 | @pytest.fixture 16 | def visitor(self): 17 | return visitors.GraphNodeVisitor() 18 | 19 | def test_has_edge(self, visitor): 20 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1')) 21 | 22 | visitor.visit(node) 23 | 24 | assert visitor.graph.get_edge(str(node), str(node.body[0])) 25 | 26 | def test_does_not_have_edge(self, visitor): 27 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1')) 28 | 29 | visitor.visit(node) 30 | 31 | assert not visitor.graph.get_edge(str(node), str(node.body[0].value)) 32 | 33 | def test_node_label(self, visitor): 34 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1')) 35 | 36 | visitor.visit(node) 37 | 38 | dot_node = visitor.graph.get_node(str(node.body[0].value))[0] 39 | if sys.version_info >= (3, 8): 40 | assert dot_node.get_label() == 'ast.Constant(value=1, kind=None)' 41 | else: 42 | assert dot_node.get_label() == 'ast.Num(n=1)' 43 | 44 | def test_edge_label(self, visitor): 45 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1')) 46 | 47 | visitor.visit(node) 48 | 49 | dot_edge = visitor.graph.get_edge(str(node), str(node.body[0]))[0] 50 | assert dot_edge.get_label() == 'body[0]' 51 | 52 | def test_multi_parents_node_label(self, visitor): 53 | node = transformers.ParentChildNodeTransformer().visit(ast.parse('x = 1\nx = 2')) 54 | 55 | visitor.visit(node) 56 | 57 | dot_node = visitor.graph.get_node(str(node.body[0].targets[0]))[0] 58 | assert dot_node.get_label() == "ast.Name(id='x', ctx=ast.Store())" 59 | 60 | 61 | class TestSourceGeneratorNodeVisitor(object): 62 | EOL = '\n' 63 | SIMPLE_ASSIGN = 'x = 1' 64 | PASS = 'pass' 65 | INDENT = ' ' * 4 66 | CLASS_DEF = 'class Sample:' 67 | EMPTY_CLASS = CLASS_DEF + EOL + INDENT + PASS 68 | FUNC_DEF = 'def f():' 69 | EMPTY_FUNC = FUNC_DEF + EOL + INDENT + PASS 70 | SINGLE_LINE_DOCSTRING = "''' This is a single line docstring.'''" 71 | MULTI_LINE_DOCSTRING = "''' This is a multi line docstring." + EOL + EOL + 'Further description...' + EOL + "'''" 72 | LINE_CONT = '\\' 73 | 74 | roundtrip_testdata = [ 75 | # assign 76 | SIMPLE_ASSIGN, 77 | '(x, y) = z', 78 | 'x += 1', 79 | 'a = b = c', 80 | '(a, b) = enumerate(c)', 81 | SIMPLE_ASSIGN + EOL + SIMPLE_ASSIGN, 82 | SIMPLE_ASSIGN + EOL + EOL + SIMPLE_ASSIGN, 83 | EOL + SIMPLE_ASSIGN, 84 | EOL + EOL + SIMPLE_ASSIGN, 85 | 'x = \'string assign\'', 86 | 87 | # class definition 88 | EMPTY_CLASS, 89 | EOL + EMPTY_CLASS, 90 | CLASS_DEF + EOL + INDENT + EOL + INDENT + PASS, 91 | EMPTY_FUNC, 92 | EOL + EMPTY_FUNC, 93 | CLASS_DEF + EOL + INDENT + FUNC_DEF + EOL + INDENT + INDENT + SIMPLE_ASSIGN, 94 | 'class A(B, C):' + EOL + INDENT + PASS, 95 | 96 | # function definition 97 | FUNC_DEF + EOL + INDENT + PASS, 98 | 'def f(x, y=1, *args, **kwargs):' + EOL + INDENT + PASS, 99 | 'def f(a, b=\'c\', *args, **kwargs):' + EOL + INDENT + PASS, 100 | FUNC_DEF + EOL + INDENT + 'return', 101 | FUNC_DEF + EOL + INDENT + 'return 5', 102 | FUNC_DEF + EOL + INDENT + 'return x == ' + LINE_CONT + EOL + INDENT + INDENT + 'x', 103 | 104 | # yield 105 | FUNC_DEF + EOL + INDENT + 'yield', 106 | FUNC_DEF + EOL + INDENT + 'yield 5', 107 | 108 | # importing 109 | 'import x', 110 | 'import x as y', 111 | 'import x.y.z', 112 | 'import x, y, z', 113 | 'from x import y', 114 | 'from x import y, z, q', 115 | 'from x import y as z', 116 | 'from x import y as z, q as p', 117 | 'from . import x', 118 | 'from .. import x', 119 | 'from .y import x', 120 | 121 | # operators 122 | '(x and y)', 123 | 'x < y', 124 | 'not x', 125 | 'x + y', 126 | '(x + y) / z', 127 | '-((-x) // y)', 128 | '(-1) ** x', 129 | '-(1 ** x)', 130 | '0 + 0j', 131 | '(-1j) ** x', 132 | 133 | # if 134 | 'if x:' + EOL + INDENT + PASS, 135 | 'if x:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + PASS, 136 | 'if x:' + EOL + INDENT + PASS + EOL + 'elif y:' + EOL + INDENT + PASS, 137 | 'if x:' + EOL + INDENT + PASS + EOL + 'elif y:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + PASS, 138 | 'if x:' + EOL + INDENT + PASS + EOL + 'elif y:' + EOL + INDENT + PASS + EOL + 'elif z:' + EOL + INDENT + PASS, 139 | 'if x:' + EOL + INDENT + PASS + EOL + 'elif y:' + EOL + INDENT + PASS + EOL + 'elif z:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + PASS, 140 | 'if x:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + 'if y:' + EOL + INDENT + INDENT + PASS + EOL + INDENT + SIMPLE_ASSIGN, 141 | 'x if y else z', 142 | 'y * (z if z > 1 else 1)', 143 | 'if x < y == z < x:' + EOL + INDENT + PASS, 144 | 'if (x < y) == (z < x):' + EOL + INDENT + PASS, 145 | 'if not False:' + EOL + INDENT + PASS, 146 | 'if x:' + EOL + INDENT + PASS + EOL + EOL + 'elif x:' + EOL + INDENT + PASS, # Double EOL 147 | 148 | # while 149 | 'while not (i != 1):' + EOL + INDENT + SIMPLE_ASSIGN, 150 | 'while True:' + EOL + INDENT + 'if True:' + EOL + INDENT + INDENT + 'continue', 151 | 'while True:' + EOL + INDENT + 'if True:' + EOL + INDENT + INDENT + 'break', 152 | SIMPLE_ASSIGN + EOL + EOL + 'while False:' + EOL + INDENT + PASS, 153 | 154 | # for 155 | 'for x in y:' + EOL + INDENT + 'break', 156 | 'for x in y:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + PASS, 157 | 158 | # try ... except 159 | 'try:' + EOL + INDENT + PASS + EOL + 'except Y:' + EOL + INDENT + PASS, 160 | 'try:' + EOL + INDENT + PASS + EOL + EOL + EOL + 'except Y:' + EOL + INDENT + PASS, 161 | 'try:' + EOL + INDENT + PASS + EOL + 'except Y as y:' + EOL + INDENT + PASS, 162 | 'try:' + EOL + INDENT + PASS + EOL + 'finally:' + EOL + INDENT + PASS, 163 | 'try:' + EOL + INDENT + PASS + EOL + 'except Y:' + EOL + INDENT + PASS + EOL + 'except Z:' + EOL + INDENT + PASS, 164 | 'try:' + EOL + INDENT + PASS + EOL + 'except Y:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + PASS, 165 | 166 | # del 167 | 'del x', 168 | 'del x, y, z', 169 | 170 | # with 171 | 'with x:' + EOL + INDENT + 'pass', 172 | 'with x as y:' + EOL + INDENT + 'pass', 173 | 174 | # assert 175 | 'assert True, \'message\'', 176 | 'assert True', 177 | 178 | # lambda 179 | 'lambda x: (x)', 180 | 'lambda x: (((x ** 2) + (2 * x)) - 5)', 181 | 'lambda: (1)', 182 | '(lambda: (yield))()', 183 | 184 | # subscript 185 | 'x[y]', 186 | 187 | # slice 188 | 'x[y:z:q]', 189 | 'x[1:2,3:4]', 190 | 'x[:2,:2]', 191 | 'x[1:2]', 192 | 'x[::2]', 193 | 194 | # global 195 | 'global x', 196 | 197 | # raise 198 | 'raise Exception()', 199 | 200 | # format 201 | '\'a %s\' % \'b\'', 202 | '\'a {}\'.format(\'b\')', 203 | '(\'%f;%f\' % (point.x, point.y)).encode(\'ascii\')', 204 | 205 | # decorator 206 | '@x(y)' + EOL + EMPTY_FUNC, 207 | 208 | # call 209 | 'f(a)', 210 | 'f(a, b)', 211 | 'f(b=\'c\')', 212 | 'f(*args)', 213 | 'f(**kwargs)', 214 | 'f(a, b=1, *args, **kwargs)', 215 | 216 | # list 217 | '[]', 218 | '[1, 2, 3]', 219 | 220 | # dict 221 | '{}', 222 | '{a: 3, b: \'c\'}', 223 | 224 | # list comprehension 225 | 'x = [y.value for y in z if y.value >= 3]', 226 | 227 | # generator expression 228 | '(x for x in y if x)', 229 | 230 | # tuple 231 | '()', 232 | '(1,)', 233 | '(1, 2)', 234 | 235 | # attribute 236 | 'x.y', 237 | 238 | # ellipsis 239 | 'x[...]', 240 | 241 | # str 242 | "x = 'y'", 243 | "x = '\"'", 244 | 'x = "\'"', 245 | 246 | # num 247 | '1', 248 | 249 | # docstring 250 | SINGLE_LINE_DOCSTRING, 251 | MULTI_LINE_DOCSTRING, 252 | CLASS_DEF + EOL + INDENT + MULTI_LINE_DOCSTRING, 253 | FUNC_DEF + EOL + INDENT + MULTI_LINE_DOCSTRING, 254 | SIMPLE_ASSIGN + EOL + MULTI_LINE_DOCSTRING, 255 | MULTI_LINE_DOCSTRING + EOL + MULTI_LINE_DOCSTRING, 256 | 257 | # line continuation 258 | 'x = ' + LINE_CONT + EOL + INDENT + 'y = 5', 259 | 'raise TypeError(' + EOL + INDENT + '\'data argument must be a bytes-like object, not str\')' 260 | ] 261 | 262 | if utils.check_version(from_inclusive=(2, 7)): 263 | roundtrip_testdata += [ 264 | # set 265 | '{1, 2}', 266 | 267 | # set comprehension 268 | '{x for x in y if x}', 269 | 270 | # dict comprehension 271 | 'x = {y: z for (y, z) in a}', 272 | ] 273 | 274 | if utils.check_version(to_exclusive=(3, 0)): 275 | roundtrip_testdata += [ 276 | # print 277 | 'print \'a\'', 278 | 'print \'a\',', 279 | 'print >> sys.stderr, \'a\'', 280 | 281 | # raise with msg and tb 282 | 'raise x, y, z', 283 | 284 | # repr 285 | '`a`', 286 | ] 287 | 288 | if utils.check_version(from_inclusive=(3, 0)): 289 | roundtrip_testdata += [ 290 | # nonlocal 291 | 'nonlocal x', 292 | 293 | # starred 294 | '*x = y', 295 | 296 | # raise from 297 | 'raise Exception() from exc', 298 | 299 | # byte string 300 | 'b\'byte_string\'', 301 | 302 | # unicode string 303 | 'x = \'äöüß\'', 304 | 305 | # metaclass 306 | 'class X(Y, metaclass=Z):' + EOL + INDENT + 'pass', 307 | 308 | # type hinting 309 | 'def f(a: str) -> str:' + EOL + INDENT + PASS, 310 | "def f(x: 'x' = 0):" + EOL + INDENT + PASS, 311 | "def f(x: 'x' = 0, *args: 'args', y: 'y' = 1, **kwargs: 'kwargs') -> 'return':" + EOL + INDENT + PASS, 312 | 313 | # extended iterable unpacking 314 | '(x, *y) = z', 315 | '[x, *y, x] = z', 316 | 317 | # kwonly arguments 318 | 'def f(*, x):' + EOL + INDENT + PASS, 319 | 'def f(*, x: int = 5):' + EOL + INDENT + PASS, 320 | 'def f(x, *, y):' + EOL + INDENT + PASS, 321 | 322 | # function definition 323 | 'def f(self, *args, x=None, **kwargs):' + EOL + INDENT + PASS, 324 | ] 325 | 326 | if utils.check_version(from_inclusive=(3, 3)): 327 | roundtrip_testdata += [ 328 | # with multiple 329 | 'with x, y:' + EOL + INDENT + 'pass', 330 | # yield from 331 | FUNC_DEF + EOL + INDENT + 'yield from x', 332 | ] 333 | 334 | if utils.check_version(from_inclusive=(3, 5)): 335 | roundtrip_testdata += [ 336 | # unpack into dict 337 | '{**kwargs}', 338 | 339 | # async/await 340 | 'async ' + FUNC_DEF + EOL + INDENT + PASS, 341 | 'async ' + FUNC_DEF + EOL + INDENT + 'async for line in reader:' + EOL + INDENT + INDENT + PASS, 342 | 'async ' + FUNC_DEF + EOL + INDENT + 'await asyncio.sleep(1)', 343 | 'async ' + FUNC_DEF + EOL + INDENT + 'async with x:' + EOL + INDENT + INDENT + PASS, 344 | 345 | # matrix multiplication operator 346 | 'x @ y', 347 | ] 348 | 349 | if utils.check_version(from_inclusive=(3, 6)): 350 | roundtrip_testdata += [ 351 | # f-strings 352 | 'f\'He said his name is {name}.\'', 353 | "f'{x!r}'", 354 | "f'{x!s}'", 355 | "f'{x!a}'", 356 | ] 357 | 358 | if utils.check_version(from_inclusive=(3, 8)): 359 | roundtrip_testdata += [ 360 | # assignment expressions 361 | 'if n := len(a) > 10:' + EOL + INDENT + PASS, 362 | # positional-only parameters 363 | 'def f(a, /, b, *, c):' + EOL + INDENT + PASS, 364 | # positional-only parameters with defaults 365 | 'def f(a=1, /, b=2, *, c=3):' + EOL + INDENT + PASS, 366 | ] 367 | 368 | # add additional tests for semantic testing 369 | semantic_testdata = list(roundtrip_testdata) 370 | 371 | semantic_testdata += [ 372 | 'x = ' + MULTI_LINE_DOCSTRING, 373 | 'b\'\'\'byte string' + EOL + 'next line' + EOL + '\'\'\'', 374 | r'r"""\a\b\f\n\r\t\v"""', 375 | 'if x:' + EOL + INDENT + PASS + EOL + 'else:' + EOL + INDENT + 'if x:' + EOL + INDENT + INDENT + PASS, 376 | ] 377 | 378 | if utils.check_version(from_inclusive=(3, 6)): 379 | semantic_testdata += [ 380 | 'raise TypeError(' + EOL + INDENT + 'f"data argument must be a bytes-like object, "' + EOL + INDENT + 381 | 'f"not {type(data).__name__}")', 382 | 'f"a\'b"', 383 | ] 384 | 385 | @pytest.mark.parametrize("source", roundtrip_testdata) 386 | def test_codegen_roundtrip(self, source): 387 | """Check if converting code into AST and converting it back to code yields the same code.""" 388 | node = ast.parse(source) 389 | generated = visitors.to_source(node) 390 | assert source == generated 391 | 392 | @pytest.mark.parametrize("source", semantic_testdata) 393 | def test_codegen_semantic_preservation(self, source): 394 | """Check if converting code into AST, converting it back to code 395 | and converting it into an AST again yields the same AST. 396 | """ 397 | node = ast.parse(source) 398 | generated = visitors.to_source(node) 399 | node_from_generated = ast.parse(generated) 400 | assert ast.dump(node) == ast.dump(node_from_generated) 401 | 402 | def test_fix_linen_umbers(self): 403 | """Check if an AST with wrong lineno attribute is corrected in the process.""" 404 | node = ast.parse('x = 1' + self.EOL + 'y = 2') 405 | # set both line numbers to 1 406 | node.body[1].lineno = 1 407 | visitors.to_source(node) 408 | assert node.body[1].lineno == 2 409 | -------------------------------------------------------------------------------- /astmonkey/transformers.py: -------------------------------------------------------------------------------- 1 | import ast 2 | 3 | 4 | class ParentChildNodeTransformer(object): 5 | 6 | def visit(self, node): 7 | self._prepare_node(node) 8 | for field, value in ast.iter_fields(node): 9 | self._process_field(node, field, value) 10 | return node 11 | 12 | @staticmethod 13 | def _prepare_node(node): 14 | if not hasattr(node, 'parent'): 15 | node.parent = None 16 | if not hasattr(node, 'parents'): 17 | node.parents = [] 18 | if not hasattr(node, 'children'): 19 | node.children = [] 20 | 21 | def _process_field(self, node, field, value): 22 | if isinstance(value, list): 23 | for index, item in enumerate(value): 24 | if isinstance(item, ast.AST): 25 | self._process_child(item, node, field, index) 26 | elif isinstance(value, ast.AST): 27 | self._process_child(value, node, field) 28 | 29 | def _process_child(self, child, parent, field_name, index=None): 30 | self.visit(child) 31 | child.parent = parent 32 | child.parents.append(parent) 33 | child.parent_field = field_name 34 | child.parent_field_index = index 35 | child.parent.children.append(child) 36 | -------------------------------------------------------------------------------- /astmonkey/utils.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import sys 3 | 4 | 5 | def is_docstring(node): 6 | if node.parent is None or node.parent.parent is None: 7 | return False 8 | def_node = node.parent.parent 9 | return ( 10 | isinstance(def_node, (ast.FunctionDef, ast.ClassDef, ast.Module)) and def_node.body and 11 | isinstance(def_node.body[0], ast.Expr) and isinstance(def_node.body[0].value, ast.Str) and 12 | def_node.body[0].value == node 13 | ) 14 | 15 | 16 | def get_by_python_version(classes, python_version=sys.version_info): 17 | result = None 18 | for cls in classes: 19 | if cls.__python_version__ <= python_version: 20 | if not result or cls.__python_version__ > result.__python_version__: 21 | result = cls 22 | if not result: 23 | raise NotImplementedError('astmonkey does not support Python %s.' % sys.version) 24 | return result 25 | 26 | 27 | class CommaWriter: 28 | 29 | def __init__(self, write_func, add_space_at_beginning=False): 30 | self.write_func = write_func 31 | self.add_space_at_beginning = add_space_at_beginning 32 | self.not_called_yet = True 33 | 34 | def __call__(self, *args, **kwargs): 35 | if self.not_called_yet: 36 | self.not_called_yet = False 37 | if self.add_space_at_beginning: 38 | self.write_func(' ') 39 | else: 40 | self.write_func(', ') 41 | 42 | def check_version(from_inclusive=None, to_exclusive=None): 43 | if (not from_inclusive or sys.version_info >= from_inclusive) \ 44 | and (not to_exclusive or sys.version_info < to_exclusive): 45 | return True 46 | return False 47 | -------------------------------------------------------------------------------- /astmonkey/visitors.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import re 3 | from contextlib import contextmanager 4 | 5 | import pydot 6 | 7 | from astmonkey import utils 8 | from astmonkey.transformers import ParentChildNodeTransformer 9 | from astmonkey.utils import CommaWriter, check_version 10 | 11 | 12 | class GraphNodeVisitor(ast.NodeVisitor): 13 | def __init__(self): 14 | self.graph = pydot.Dot(graph_type='graph', **self._dot_graph_kwargs()) 15 | 16 | def visit(self, node): 17 | if len(node.parents) <= 1: 18 | self.graph.add_node(self._dot_node(node)) 19 | if len(node.parents) == 1: 20 | self.graph.add_edge(self._dot_edge(node)) 21 | super(GraphNodeVisitor, self).visit(node) 22 | 23 | def _dot_graph_kwargs(self): 24 | return {} 25 | 26 | def _dot_node(self, node): 27 | return pydot.Node(str(node), label=self._dot_node_label(node), **self._dot_node_kwargs(node)) 28 | 29 | def _dot_node_label(self, node): 30 | fields_labels = [] 31 | for field, value in ast.iter_fields(node): 32 | if not isinstance(value, list): 33 | value_label = self._dot_node_value_label(value) 34 | if value_label: 35 | fields_labels.append('{0}={1}'.format(field, value_label)) 36 | return 'ast.{0}({1})'.format(node.__class__.__name__, ', '.join(fields_labels)) 37 | 38 | def _dot_node_value_label(self, value): 39 | if not isinstance(value, ast.AST): 40 | return repr(value) 41 | elif len(value.parents) > 1: 42 | return self._dot_node_label(value) 43 | return None 44 | 45 | def _dot_node_kwargs(self, node): 46 | return { 47 | 'shape': 'box', 48 | 'fontname': 'Curier' 49 | } 50 | 51 | def _dot_edge(self, node): 52 | return pydot.Edge(str(node.parent), str(node), label=self._dot_edge_label(node), **self._dot_edge_kwargs(node)) 53 | 54 | def _dot_edge_label(self, node): 55 | label = node.parent_field 56 | if not node.parent_field_index is None: 57 | label += '[{0}]'.format(node.parent_field_index) 58 | return label 59 | 60 | def _dot_edge_kwargs(self, node): 61 | return { 62 | 'fontname': 'Curier' 63 | } 64 | 65 | 66 | """ 67 | Source generator node visitor from Python AST was originaly written by Armin Ronacher (2008), license BSD. 68 | """ 69 | 70 | BOOLOP_SYMBOLS = { 71 | ast.And: 'and', 72 | ast.Or: 'or' 73 | } 74 | 75 | BINOP_SYMBOLS = { 76 | ast.Add: '+', 77 | ast.Sub: '-', 78 | ast.Mult: '*', 79 | ast.Div: '/', 80 | ast.FloorDiv: '//', 81 | ast.Mod: '%', 82 | ast.LShift: '<<', 83 | ast.RShift: '>>', 84 | ast.BitOr: '|', 85 | ast.BitAnd: '&', 86 | ast.BitXor: '^', 87 | ast.Pow: '**' 88 | } 89 | 90 | if check_version(from_inclusive=(3, 5)): 91 | BINOP_SYMBOLS[ast.MatMult] = '@' 92 | 93 | CMPOP_SYMBOLS = { 94 | ast.Eq: '==', 95 | ast.Gt: '>', 96 | ast.GtE: '>=', 97 | ast.In: 'in', 98 | ast.Is: 'is', 99 | ast.IsNot: 'is not', 100 | ast.Lt: '<', 101 | ast.LtE: '<=', 102 | ast.NotEq: '!=', 103 | ast.NotIn: 'not in' 104 | } 105 | 106 | UNARYOP_SYMBOLS = { 107 | ast.Invert: '~', 108 | ast.Not: 'not', 109 | ast.UAdd: '+', 110 | ast.USub: '-' 111 | } 112 | 113 | ALL_SYMBOLS = {} 114 | ALL_SYMBOLS.update(BOOLOP_SYMBOLS) 115 | ALL_SYMBOLS.update(BINOP_SYMBOLS) 116 | ALL_SYMBOLS.update(CMPOP_SYMBOLS) 117 | ALL_SYMBOLS.update(UNARYOP_SYMBOLS) 118 | 119 | 120 | def to_source(node, indent_with=' ' * 4): 121 | """This function can convert a node tree back into python sourcecode. 122 | This is useful for debugging purposes, especially if you're dealing with 123 | custom asts not generated by python itself. 124 | 125 | It could be that the sourcecode is evaluable when the AST itself is not 126 | compilable / evaluable. The reason for this is that the AST contains some 127 | more data than regular sourcecode does, which is dropped during 128 | conversion. 129 | 130 | Each level of indentation is replaced with `indent_with`. Per default this 131 | parameter is equal to four spaces as suggested by PEP 8, but it might be 132 | adjusted to match the application's styleguide. 133 | """ 134 | ParentChildNodeTransformer().visit(node) 135 | FixLinenoNodeVisitor().visit(node) 136 | generator = SourceGeneratorNodeVisitor(indent_with) 137 | generator.visit(node) 138 | 139 | return ''.join(generator.result) 140 | 141 | 142 | class FixLinenoNodeVisitor(ast.NodeVisitor): 143 | """A helper node visitor for the SourceGeneratorNodeVisitor. 144 | 145 | Attempts to correct implausible line numbers. An example would be: 146 | 147 | 1: while a: 148 | 2: pass 149 | 3: for a: 150 | 2: pass 151 | 152 | This would be corrected to: 153 | 154 | 1: while a: 155 | 2: pass 156 | 3: for a: 157 | 4: pass 158 | """ 159 | 160 | def __init__(self): 161 | self.min_lineno = 0 162 | 163 | def generic_visit(self, node): 164 | if hasattr(node, 'lineno'): 165 | self._fix_lineno(node) 166 | if hasattr(node, 'body') and isinstance(node.body, list): 167 | self._process_body(node) 168 | 169 | def _fix_lineno(self, node): 170 | if node.lineno < self.min_lineno: 171 | node.lineno = self.min_lineno 172 | else: 173 | self.min_lineno = node.lineno 174 | 175 | def _process_body(self, node): 176 | for body_node in node.body: 177 | self.min_lineno += 1 178 | self.visit(body_node) 179 | 180 | 181 | class BaseSourceGeneratorNodeVisitor(ast.NodeVisitor): 182 | """This visitor is able to transform a well formed syntax tree into python 183 | sourcecode. For more details have a look at the docstring of the 184 | `node_to_source` function. 185 | """ 186 | 187 | def __init__(self, indent_with): 188 | self.result = [] 189 | self.indent_with = indent_with 190 | self.indentation = 0 191 | 192 | @classmethod 193 | def _is_node_args_valid(cls, node, arg_name): 194 | return hasattr(node, arg_name) and getattr(node, arg_name) is not None 195 | 196 | def _get_current_line_no(self): 197 | lines = len("".join(self.result).split('\n')) if self.result else 0 198 | return lines 199 | 200 | @classmethod 201 | def _get_actual_lineno(cls, node): 202 | if isinstance(node, (ast.Expr, ast.Str)) and node.col_offset == -1: 203 | str_content = cls._get_string_content(node) 204 | node_lineno = node.lineno - str_content.count('\n') 205 | else: 206 | node_lineno = node.lineno 207 | return node_lineno 208 | 209 | @staticmethod 210 | def _get_string_content(node): 211 | # node is a multi line string and the line number is actually the last line 212 | if isinstance(node, ast.Expr): 213 | str_content = node.value.s 214 | else: 215 | str_content = node.s 216 | if type(str_content) == bytes: 217 | str_content = str_content.decode("utf-8") 218 | return str_content 219 | 220 | def _newline_needed(self, node): 221 | lines = self._get_current_line_no() 222 | node_lineno = self._get_actual_lineno(node) 223 | line_diff = node_lineno - lines 224 | return line_diff > 0 225 | 226 | @contextmanager 227 | def indent(self, count=1): 228 | self.indentation += count 229 | yield 230 | self.indentation -= count 231 | 232 | @contextmanager 233 | def inside(self, pre, post, cond=True): 234 | if cond: 235 | self.write(pre) 236 | yield 237 | if cond: 238 | self.write(post) 239 | 240 | def write(self, x): 241 | self.result.append(x) 242 | 243 | def correct_line_number(self, node, within_statement=True, use_line_continuation=True): 244 | if not node or not self._is_node_args_valid(node, 'lineno'): 245 | return 246 | if within_statement: 247 | indent = 1 248 | else: 249 | indent = 0 250 | with self.indent(indent): 251 | self.add_missing_lines(node, within_statement, use_line_continuation) 252 | 253 | def add_missing_lines(self, node, within_statement, use_line_continuation): 254 | while self._newline_needed(node): 255 | self.add_line(within_statement, use_line_continuation) 256 | 257 | def add_line(self, within_statement, use_line_continuation): 258 | if within_statement and use_line_continuation: 259 | self.result.append('\\') 260 | self.write_newline() 261 | 262 | def write_newline(self): 263 | if self.result: 264 | self.result.append('\n') 265 | self.result.append(self.indent_with * self.indentation) 266 | 267 | def body(self, statements, indent=1): 268 | if statements: 269 | with self.indent(indent): 270 | for stmt in statements: 271 | self.correct_line_number(stmt, within_statement=False) 272 | self.visit(stmt) 273 | 274 | def body_or_else(self, node): 275 | self.body(node.body) 276 | if node.orelse: 277 | self.or_else(node) 278 | 279 | def keyword_and_body(self, keyword, body): 280 | if self._newline_needed(body[0]): 281 | self.write_newline() 282 | self.write(keyword) 283 | self.body(body) 284 | 285 | def or_else(self, node): 286 | self.keyword_and_body('else:', node.orelse) 287 | 288 | def docstring(self, node): 289 | s = repr(node.s) 290 | s = re.sub(r'(?> ') 484 | self.visit(node.dest) 485 | want_comma = True 486 | for value in node.values: 487 | if want_comma: 488 | self.write(', ') 489 | self.visit(value) 490 | want_comma = True 491 | if not node.nl: 492 | self.write(',') 493 | 494 | def visit_Delete(self, node): 495 | self.write('del ') 496 | for target in node.targets: 497 | self.visit(target) 498 | if target is not node.targets[-1]: 499 | self.write(', ') 500 | 501 | def visit_Global(self, node): 502 | self.write('global ' + ', '.join(node.names)) 503 | 504 | def visit_Nonlocal(self, node): 505 | self.write('nonlocal ' + ', '.join(node.names)) 506 | 507 | def visit_Return(self, node): 508 | 509 | self.write('return') 510 | if node.value: 511 | self.write(' ') 512 | self.visit(node.value) 513 | 514 | def visit_Break(self, node): 515 | self.write('break') 516 | 517 | def visit_Continue(self, node): 518 | self.write('continue') 519 | 520 | def visit_Raise(self, node): 521 | 522 | self.write('raise') 523 | if self._is_node_args_valid(node, 'exc'): 524 | self.raise_exc(node) 525 | elif self._is_node_args_valid(node, 'type'): 526 | self.raise_type(node) 527 | 528 | def raise_type(self, node): 529 | self.write(' ') 530 | self.visit(node.type) 531 | if node.inst is not None: 532 | self.write(', ') 533 | self.visit(node.inst) 534 | if node.tback is not None: 535 | self.write(', ') 536 | self.visit(node.tback) 537 | 538 | def raise_exc(self, node): 539 | self.write(' ') 540 | self.visit(node.exc) 541 | if node.cause is not None: 542 | self.write(' from ') 543 | self.visit(node.cause) 544 | 545 | # Expressions 546 | 547 | def visit_Attribute(self, node): 548 | self.visit(node.value) 549 | self.write('.' + node.attr) 550 | 551 | def visit_Call(self, node): 552 | self.visit(node.func) 553 | with self.inside('(', ')'): 554 | starargs = getattr(node, 'starargs', None) 555 | kwargs = getattr(node, 'kwargs', None) 556 | if starargs: 557 | starargs = [starargs] 558 | else: 559 | starargs = [] 560 | if kwargs: 561 | kwargs = [kwargs] 562 | else: 563 | kwargs = [] 564 | self.call_signature(node.args, node.keywords, starargs, kwargs) 565 | 566 | def call_signature(self, args, keywords, starargs, kwargs): 567 | write_comma = CommaWriter(self.write) 568 | self.call_signature_part(args, self.call_arg, write_comma) 569 | self.call_signature_part(keywords, self.call_keyword, write_comma) 570 | self.call_signature_part(starargs, self.call_starargs, write_comma) 571 | self.call_signature_part(kwargs, self.call_kwarg, write_comma) 572 | 573 | def call_signature_part(self, args, arg_processor, write_comma): 574 | for arg in args: 575 | write_comma() 576 | self.correct_line_number(arg, use_line_continuation=False) 577 | arg_processor(arg) 578 | 579 | def call_kwarg(self, kwarg): 580 | self.write('**') 581 | self.visit(kwarg) 582 | 583 | def call_starargs(self, stararg): 584 | self.write('*') 585 | self.visit(stararg) 586 | 587 | def call_keyword(self, keyword): 588 | self.visit(keyword) 589 | 590 | def call_arg(self, arg): 591 | self.visit(arg) 592 | 593 | def visit_Name(self, node): 594 | self.write(node.id) 595 | 596 | def visit_str(self, node): 597 | self.write(node) 598 | 599 | def visit_Str(self, node): 600 | self.write(repr(node.s)) 601 | 602 | def visit_Bytes(self, node): 603 | self.write(repr(node.s)) 604 | 605 | def visit_Num(self, node): 606 | value = node.n.imag if isinstance(node.n, complex) else node.n 607 | 608 | with self.inside('(', ')', cond=(value < 0)): 609 | self.write(repr(node.n)) 610 | 611 | def visit_Tuple(self, node): 612 | with self.inside('(', ')'): 613 | idx = -1 614 | for idx, item in enumerate(node.elts): 615 | if idx: 616 | self.write(', ') 617 | self.visit(item) 618 | if not idx: 619 | self.write(',') 620 | 621 | def sequence_visit(left, right): # @NoSelf 622 | def visit(self, node): 623 | with self.inside(left, right): 624 | for idx, item in enumerate(node.elts): 625 | if idx: 626 | self.write(', ') 627 | self.visit(item) 628 | 629 | return visit 630 | 631 | visit_List = sequence_visit('[', ']') 632 | visit_Set = sequence_visit('{', '}') 633 | del sequence_visit 634 | 635 | def visit_Dict(self, node): 636 | with self.inside('{', '}'): 637 | for idx, (key, value) in enumerate(zip(node.keys, node.values)): 638 | if idx: 639 | self.write(', ') 640 | if key: 641 | self.visit(key) 642 | self.write(': ') 643 | else: 644 | self.write('**') 645 | self.visit(value) 646 | 647 | def visit_BinOp(self, node): 648 | with self.inside('(', ')', cond=isinstance(node.parent, (ast.BinOp, ast.Attribute))): 649 | self.visit(node.left) 650 | self.write(' %s ' % BINOP_SYMBOLS[type(node.op)]) 651 | self.visit(node.right) 652 | 653 | def visit_BoolOp(self, node): 654 | with self.inside('(', ')'): 655 | for idx, value in enumerate(node.values): 656 | if idx: 657 | self.write(' %s ' % BOOLOP_SYMBOLS[type(node.op)]) 658 | self.visit(value) 659 | 660 | def visit_Compare(self, node): 661 | with self.inside('(', ')', cond=(isinstance(node.parent, ast.Compare))): 662 | self.visit(node.left) 663 | for op, right in zip(node.ops, node.comparators): 664 | self.write(' %s ' % CMPOP_SYMBOLS[type(op)]) 665 | self.visit(right) 666 | 667 | def visit_UnaryOp(self, node): 668 | with self.inside('(', ')', cond=isinstance(node.parent, (ast.BinOp, ast.UnaryOp))): 669 | op = UNARYOP_SYMBOLS[type(node.op)] 670 | self.write(op) 671 | if op == 'not': 672 | self.write(' ') 673 | 674 | with self.inside('(', ')', cond=(not isinstance(node.operand, (ast.Name, ast.Num)) 675 | and not self._is_named_constant(node.operand))): 676 | self.visit(node.operand) 677 | 678 | def visit_Subscript(self, node): 679 | self.visit(node.value) 680 | with self.inside('[', ']'): 681 | self.visit(node.slice) 682 | 683 | def visit_Slice(self, node): 684 | self.slice_lower(node) 685 | self.write(':') 686 | self.slice_upper(node) 687 | self.slice_step(node) 688 | 689 | def slice_step(self, node): 690 | if node.step is not None: 691 | self.write(':') 692 | if not (isinstance(node.step, ast.Name) and node.step.id == 'None'): 693 | self.visit(node.step) 694 | 695 | def slice_upper(self, node): 696 | if node.upper is not None: 697 | self.visit(node.upper) 698 | 699 | def slice_lower(self, node): 700 | if node.lower is not None: 701 | self.visit(node.lower) 702 | 703 | def visit_ExtSlice(self, node): 704 | for idx, item in enumerate(node.dims): 705 | if idx: 706 | self.write(',') 707 | self.visit(item) 708 | 709 | def visit_Yield(self, node): 710 | self.write('yield') 711 | if node.value: 712 | self.write(' ') 713 | self.visit(node.value) 714 | 715 | def visit_Lambda(self, node): 716 | with self.inside('(', ')', cond=isinstance(node.parent, ast.Call)): 717 | self.write('lambda') 718 | self.signature(node.args, add_space=True) 719 | self.write(': ') 720 | with self.inside('(', ')'): 721 | self.visit(node.body) 722 | 723 | def visit_Ellipsis(self, node): 724 | self.write('...') 725 | 726 | def generator_visit(left, right): # @NoSelf 727 | def visit(self, node): 728 | self.write(left) 729 | self.visit(node.elt) 730 | for comprehension in node.generators: 731 | self.visit(comprehension) 732 | self.write(right) 733 | 734 | return visit 735 | 736 | visit_ListComp = generator_visit('[', ']') 737 | visit_GeneratorExp = generator_visit('(', ')') 738 | visit_SetComp = generator_visit('{', '}') 739 | del generator_visit 740 | 741 | def visit_DictComp(self, node): 742 | with self.inside('{', '}'): 743 | self.visit(node.key) 744 | self.write(': ') 745 | self.visit(node.value) 746 | for comprehension in node.generators: 747 | self.visit(comprehension) 748 | 749 | def visit_IfExp(self, node): 750 | with self.inside('(', ')', cond=isinstance(node.parent, ast.BinOp)): 751 | self.visit(node.body) 752 | self.write(' if ') 753 | self.visit(node.test) 754 | self.keyword_and_body(' else ', [node.orelse]) 755 | 756 | def visit_Starred(self, node): 757 | self.write('*') 758 | self.visit(node.value) 759 | 760 | def visit_Repr(self, node): 761 | with self.inside('`', '`'): 762 | self.visit(node.value) 763 | 764 | # Helper Nodes 765 | def visit_alias(self, node): 766 | self.write(node.name) 767 | if node.asname is not None: 768 | self.write(' as ' + node.asname) 769 | 770 | def visit_comprehension(self, node): 771 | self.write(' for ') 772 | self.visit(node.target) 773 | self.write(' in ') 774 | self.visit(node.iter) 775 | if node.ifs: 776 | for if_ in node.ifs: 777 | self.write(' if ') 778 | self.visit(if_) 779 | 780 | def visit_ExceptHandler(self, node): 781 | self.write('except') 782 | if node.type is not None: 783 | self.write(' ') 784 | self.visit(node.type) 785 | if node.name is not None: 786 | self.write(' as ') 787 | self.visit(node.name) 788 | self.write(':') 789 | self.body(node.body) 790 | 791 | def visit_arg(self, node): 792 | self.write(node.arg) 793 | 794 | def visit_Assert(self, node): 795 | self.write('assert ') 796 | self.visit(node.test) 797 | if node.msg: 798 | self.write(', ') 799 | self.visit(node.msg) 800 | 801 | def visit_TryExcept(self, node): 802 | self.write('try:') 803 | self.body(node.body) 804 | if node.handlers: 805 | self.try_handlers(node) 806 | if node.orelse: 807 | self.or_else(node) 808 | 809 | def try_handlers(self, node): 810 | for handler in node.handlers: 811 | self.correct_line_number(handler, within_statement=False) 812 | self.visit(handler) 813 | 814 | def visit_TryFinally(self, node): 815 | self.write('try:') 816 | self.body(node.body) 817 | self.final_body(node) 818 | 819 | def final_body(self, node): 820 | self.keyword_and_body('finally:', node.finalbody) 821 | 822 | def visit_With(self, node): 823 | self.with_body(node) 824 | 825 | def with_body(self, node, prefixes=[]): 826 | self._prefixes(prefixes) 827 | self.write('with ') 828 | self.visit(node.context_expr) 829 | if node.optional_vars is not None: 830 | self.write(' as ') 831 | self.visit(node.optional_vars) 832 | self.write(':') 833 | self.body(node.body) 834 | 835 | @staticmethod 836 | def _is_named_constant(node): 837 | return isinstance(node, ast.Expr) and hasattr(node, 'value') and isinstance(node.value, ast.Name) 838 | 839 | 840 | class SourceGeneratorNodeVisitorPython26(BaseSourceGeneratorNodeVisitor): 841 | __python_version__ = (2, 6) 842 | 843 | 844 | class SourceGeneratorNodeVisitorPython27(SourceGeneratorNodeVisitorPython26): 845 | __python_version__ = (2, 7) 846 | 847 | 848 | class SourceGeneratorNodeVisitorPython30(SourceGeneratorNodeVisitorPython27): 849 | __python_version__ = (3, 0) 850 | 851 | def visit_ClassDef(self, node): 852 | have_args = [] 853 | 854 | def paren_or_comma(): 855 | if have_args: 856 | self.write(', ') 857 | else: 858 | have_args.append(True) 859 | self.write('(') 860 | 861 | self.decorators(node) 862 | self.correct_line_number(node) 863 | self.write('class %s' % node.name) 864 | for base in node.bases: 865 | paren_or_comma() 866 | self.visit(base) 867 | if self._is_node_args_valid(node, 'keywords'): 868 | for keyword in node.keywords: 869 | paren_or_comma() 870 | self.visit(keyword) 871 | self.write(have_args and '):' or ':') 872 | self.body(node.body) 873 | 874 | def visit_FunctionDef(self, node): 875 | self.decorators(node) 876 | 877 | self.write('def %s(' % node.name) 878 | self.signature(node.args) 879 | self.write(')') 880 | if self._is_node_args_valid(node, 'returns'): 881 | self.write(' -> ') 882 | self.visit(node.returns) 883 | self.write(':') 884 | self.body(node.body) 885 | 886 | 887 | class SourceGeneratorNodeVisitorPython31(SourceGeneratorNodeVisitorPython30): 888 | __python_version__ = (3, 1) 889 | 890 | 891 | class SourceGeneratorNodeVisitorPython32(SourceGeneratorNodeVisitorPython31): 892 | __python_version__ = (3, 2) 893 | 894 | 895 | class SourceGeneratorNodeVisitorPython33(SourceGeneratorNodeVisitorPython32): 896 | __python_version__ = (3, 3) 897 | 898 | def visit_Try(self, node): 899 | self.write('try:') 900 | self.body(node.body) 901 | if node.handlers: 902 | self.try_handlers(node) 903 | if node.finalbody: 904 | self.final_body(node) 905 | if node.orelse: 906 | self.or_else(node) 907 | 908 | def with_body(self, node, prefixes=[]): 909 | self._prefixes(prefixes) 910 | self.write('with ') 911 | for with_item in node.items: 912 | self.visit(with_item.context_expr) 913 | if with_item.optional_vars is not None: 914 | self.write(' as ') 915 | self.visit(with_item.optional_vars) 916 | if with_item != node.items[-1]: 917 | self.write(', ') 918 | self.write(':') 919 | self.body(node.body) 920 | 921 | def visit_YieldFrom(self, node): 922 | self.write('yield from ') 923 | self.visit(node.value) 924 | 925 | 926 | class SourceGeneratorNodeVisitorPython34(SourceGeneratorNodeVisitorPython33): 927 | __python_version__ = (3, 4) 928 | 929 | def visit_NameConstant(self, node): 930 | self.write(str(node.value)) 931 | 932 | def visit_Name(self, node): 933 | if isinstance(node.id, ast.arg): 934 | self.write(node.id.arg) 935 | else: 936 | self.write(node.id) 937 | 938 | @staticmethod 939 | def _is_named_constant(node): 940 | return isinstance(node, ast.NameConstant) 941 | 942 | 943 | class SourceGeneratorNodeVisitorPython35(SourceGeneratorNodeVisitorPython34): 944 | __python_version__ = (3, 5) 945 | 946 | def visit_AsyncFunctionDef(self, node): 947 | self.function_definition(node, prefixes=['async']) 948 | 949 | def visit_AsyncFor(self, node): 950 | self.for_loop(node, prefixes=['async']) 951 | 952 | def visit_AsyncWith(self, node): 953 | self.with_body(node, prefixes=['async']) 954 | 955 | def visit_Await(self, node): 956 | self.write('await ') 957 | if self._is_node_args_valid(node, 'value'): 958 | self.visit(node.value) 959 | 960 | def visit_Call(self, node): 961 | self.visit(node.func) 962 | with self.inside('(', ')'): 963 | args, starargs = self._separate_args_and_starargs(node) 964 | keywords, kwargs = self._separate_keywords_and_kwargs(node) 965 | self.call_signature(args, keywords, starargs, kwargs) 966 | 967 | @staticmethod 968 | def _separate_keywords_and_kwargs(node): 969 | keywords = [] 970 | kwargs = [] 971 | for keyword in node.keywords: 972 | if keyword.arg: 973 | keywords.append(keyword) 974 | else: 975 | kwargs.append(keyword) 976 | return keywords, kwargs 977 | 978 | @staticmethod 979 | def _separate_args_and_starargs(node): 980 | args = [] 981 | starargs = [] 982 | for arg in node.args: 983 | if isinstance(arg, ast.Starred): 984 | starargs.append(arg) 985 | else: 986 | args.append(arg) 987 | return args, starargs 988 | 989 | def call_starargs(self, stararg): 990 | self.visit(stararg) 991 | 992 | def call_kwarg(self, kwarg): 993 | self.visit(kwarg) 994 | 995 | 996 | class SourceGeneratorNodeVisitorPython36(SourceGeneratorNodeVisitorPython35): 997 | __python_version__ = (3, 6) 998 | 999 | def visit_JoinedStr(self, node): 1000 | if self._is_node_args_valid(node, 'values'): 1001 | with self.inside('f\'', '\''): 1002 | for item in node.values: 1003 | if isinstance(item, ast.Str): 1004 | self.write(item.s.lstrip('\'').rstrip('\'').replace("'", "\\'")) 1005 | else: 1006 | self.visit(item) 1007 | 1008 | def visit_FormattedValue(self, node): 1009 | if self._is_node_args_valid(node, 'value'): 1010 | with self.inside('{', '}'): 1011 | self.visit(node.value) 1012 | if node.conversion != -1: 1013 | self.write('!%c' % (node.conversion,)) 1014 | 1015 | 1016 | class SourceGeneratorNodeVisitorPython38(SourceGeneratorNodeVisitorPython36): 1017 | __python_version__ = (3, 8) 1018 | 1019 | def visit_Constant(self, node): 1020 | if type(node.value) == str: 1021 | self.write(repr(node.s)) 1022 | elif node.value == Ellipsis: 1023 | self.write('...') 1024 | else: 1025 | self.write(str(node.value)) 1026 | 1027 | def visit_NamedExpr(self, node): 1028 | self.visit(node.target) 1029 | self.write(' := ') 1030 | self.visit(node.value) 1031 | 1032 | def signature(self, node, add_space=False): 1033 | write_comma = CommaWriter(self.write, add_space_at_beginning=add_space) 1034 | 1035 | 1036 | defaults = list(node.defaults) 1037 | 1038 | if node.posonlyargs: 1039 | padding = [None] * (len(node.posonlyargs) - len(node.defaults)) 1040 | for arg, default in zip(node.posonlyargs, padding + defaults[:len(node.posonlyargs)]): 1041 | self.signature_arg(arg, default, write_comma) 1042 | self.write(', /') 1043 | defaults = defaults[len(node.posonlyargs):] 1044 | 1045 | padding = [None] * (len(node.args) - len(node.defaults)) 1046 | for arg, default in zip(node.args, padding + defaults): 1047 | self.signature_arg(arg, default, write_comma) 1048 | 1049 | self.signature_spec_arg(node, 'vararg', write_comma, prefix='*') 1050 | self.signature_kwonlyargs(node, write_comma) 1051 | self.signature_spec_arg(node, 'kwarg', write_comma, prefix='**') 1052 | 1053 | @classmethod 1054 | def _get_actual_lineno(cls, node): 1055 | if isinstance(node, ast.FunctionDef) and node.decorator_list: 1056 | return node.decorator_list[0].lineno 1057 | else: 1058 | return SourceGeneratorNodeVisitorPython36._get_actual_lineno(node) 1059 | 1060 | 1061 | SourceGeneratorNodeVisitor = utils.get_by_python_version([ 1062 | SourceGeneratorNodeVisitorPython26, 1063 | SourceGeneratorNodeVisitorPython27, 1064 | SourceGeneratorNodeVisitorPython30, 1065 | SourceGeneratorNodeVisitorPython31, 1066 | SourceGeneratorNodeVisitorPython32, 1067 | SourceGeneratorNodeVisitorPython33, 1068 | SourceGeneratorNodeVisitorPython34, 1069 | SourceGeneratorNodeVisitorPython35, 1070 | SourceGeneratorNodeVisitorPython36, 1071 | SourceGeneratorNodeVisitorPython38 1072 | ]) 1073 | -------------------------------------------------------------------------------- /examples/edge-graph-node-visitor.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """This example was kindly provided by https://github.com/oozie. 3 | 4 | This script draws the AST of a Python module as graph with simple points as nodes. 5 | Astmonkey's GraphNodeVisitor is subclassed for custom representation of the AST. 6 | 7 | Usage: python3 edge-graph-node-visitor.py some_file.py 8 | """ 9 | import ast 10 | import os 11 | import sys 12 | 13 | import pydot 14 | 15 | from astmonkey import transformers 16 | from astmonkey.visitors import GraphNodeVisitor 17 | 18 | 19 | class EdgeGraphNodeVisitor(GraphNodeVisitor): 20 | """Simple point-edge-point graphviz representation of the AST.""" 21 | 22 | def __init__(self): 23 | super(self.__class__, self).__init__() 24 | self.graph.set_node_defaults(shape='point') 25 | 26 | def _dot_graph_kwargs(self): 27 | return {} 28 | 29 | def _dot_node_kwargs(self, node): 30 | return {} 31 | 32 | def _dot_edge(self, node): 33 | return pydot.Edge(id(node.parent), id(node)) 34 | 35 | def _dot_node(self, node): 36 | return pydot.Node(id(node), **self._dot_node_kwargs(node)) 37 | 38 | 39 | if __name__ == '__main__': 40 | filename = sys.argv[1] 41 | 42 | node = ast.parse(open(filename).read()) 43 | node = transformers.ParentChildNodeTransformer().visit(node) 44 | visitor = EdgeGraphNodeVisitor() 45 | visitor.visit(node) 46 | visitor.graph.write(filename + '.dot') 47 | os.system('sfdp -Tpng -o {} {}'.format(filename + '.png', filename + '.dot')) 48 | -------------------------------------------------------------------------------- /examples/edge-graph.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mutpy/astmonkey/16b1817fd022b60b1369281af74b465db6db85ff/examples/edge-graph.png -------------------------------------------------------------------------------- /examples/graph.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mutpy/astmonkey/16b1817fd022b60b1369281af74b465db6db85ff/examples/graph.png -------------------------------------------------------------------------------- /examples/graph_node_visitor.py: -------------------------------------------------------------------------------- 1 | import ast 2 | from astmonkey import visitors, transformers 3 | 4 | node = ast.parse('def foo(x):\n\treturn x + 1') 5 | node = transformers.ParentChildNodeTransformer().visit(node) 6 | visitor = visitors.GraphNodeVisitor() 7 | visitor.visit(node) 8 | 9 | visitor.graph.write_png('graph.png') 10 | -------------------------------------------------------------------------------- /examples/is_docstring.py: -------------------------------------------------------------------------------- 1 | import ast 2 | from astmonkey import utils, transformers 3 | 4 | node = ast.parse('def foo(x):\n\t"""doc"""') 5 | node = transformers.ParentChildNodeTransformer().visit(node) 6 | 7 | docstring_node = node.body[0].body[0].value 8 | assert(not utils.is_docstring(node)) 9 | assert(utils.is_docstring(docstring_node)) 10 | -------------------------------------------------------------------------------- /examples/parent_node_transformer.py: -------------------------------------------------------------------------------- 1 | import ast 2 | from astmonkey import transformers 3 | 4 | node = ast.parse('x = 1') 5 | node = transformers.ParentChildNodeTransformer().visit(node) 6 | 7 | assert(node == node.body[0].parent) 8 | assert(node.body[0].parent_field == 'body') 9 | assert(node.body[0].parent_field_index == 0) 10 | -------------------------------------------------------------------------------- /examples/source_generator_node_visitor.py: -------------------------------------------------------------------------------- 1 | import ast 2 | from astmonkey import visitors 3 | 4 | code = 'x = y + 1' 5 | node = ast.parse(code) 6 | generated_code = visitors.to_source(node) 7 | 8 | assert(code == generated_code) 9 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from setuptools import setup 3 | 4 | import astmonkey 5 | 6 | with open('README.rst') as f: 7 | long_description = f.read() 8 | 9 | setup( 10 | name='astmonkey', 11 | version=astmonkey.__version__, 12 | description='astmonkey is a set of tools to play with Python AST.', 13 | author='Konrad Hałas', 14 | author_email='halas.konrad@gmail.com', 15 | url='https://github.com/mutpy/astmonkey', 16 | packages=['astmonkey'], 17 | install_requires=['pydot'], 18 | long_description=long_description, 19 | classifiers=[ 20 | 'Intended Audience :: Developers', 21 | 'Programming Language :: Python :: 2.6', 22 | 'Programming Language :: Python :: 2.7', 23 | 'Programming Language :: Python :: 3.3', 24 | 'Programming Language :: Python :: 3.4', 25 | 'Programming Language :: Python :: 3.5', 26 | 'Programming Language :: Python :: 3.6', 27 | 'Programming Language :: Python :: 3.7', 28 | 'Programming Language :: Python :: 3.8', 29 | 'License :: OSI Approved :: Apache Software License' 30 | ] 31 | ) 32 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = 3 | coverage-erase 4 | test-py{27,34,35,36,37,38} 5 | coverage-report 6 | [testenv] 7 | deps= 8 | coverage 9 | test-py26: unittest2 10 | test: pydot 11 | test: pytest 12 | test: pytest-cov 13 | commands = 14 | coverage-erase: coverage erase 15 | test: py.test --cov=astmonkey --cov-report= --cov-append astmonkey/tests 16 | coverage-report: coverage report 17 | --------------------------------------------------------------------------------