├── LICENSE
├── README.md
├── astree.py
└── requirements.txt


/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2020 Ira Horecka
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # astree
 2 | ### Visualize abstract syntax trees of methods, declarations, expressions, and more
 3 | 
 4 | ```astree.py``` combines the Python <a href=https://docs.python.org/3.8/library/ast.html>```ast```</a> module with
 5 | <a href=https://github.com/pydot/pydot-ng>```pydot_ng```</a> to draw abstract syntax trees specified in DOT language scripts.
 6 | An <a href=https://en.wikipedia.org/wiki/Abstract_syntax_tree> abstract syntax tree</a> is a tree representation of the abstract
 7 | syntactic structure of source code written in a programming language (e.g. Python).<br>
 8 | <hr>
 9 |   
10 | <b>Jumpstart</b> -- running the program:
11 | 1) Clone repository
12 | 2) ```$ pip install -r requirements.txt```<br>
13 | 3) ```$ python astree.py```
14 | 
15 | Input modules, methods, declarations, statements, expressions, etc.<br>
16 | View video example <a href="https://i.imgur.com/C43Mn6p.mp4">here</a>.<br>
17 | 
18 | For example, let's look at the ```requests.get``` method:<br>
19 | ```>>> Input a method name, expression, etc.:```<br>
20 | ```requests.get```<br>
21 | 
22 | <p align = 'center'>
23 | <img src=https://i.imgur.com/41FcAwg.png alt="AST visualize requests.get"
24 |     width=800><br>
25 | </p>
26 | <br>
27 | Note: please report bugs to <a href="https://github.com/irahorecka/astree/issues">issues</a>.
28 | 


--------------------------------------------------------------------------------
/astree.py:
--------------------------------------------------------------------------------
  1 | """A module to visualize Python AST."""
  2 | 
  3 | import ast
  4 | import inspect
  5 | import importlib
  6 | import json
  7 | import re
  8 | import uuid
  9 | import pydot_ng as pydot
 10 | from IPython.display import Image, display
 11 | from _ast import AST
 12 | 
 13 | 
 14 | # ~~~~~~~~~~~~~~~~~~ PARSING AST OBJ TO JSON ~~~~~~~~~~~~~~~~~~
 15 | 
 16 | 
 17 | def ast_parse(method):
 18 |     """Decorator to parse user input to JSON-AST object."""
 19 | 
 20 |     def wrapper(*args, **kwargs):
 21 |         if isinstance(args[0], str):
 22 |             ast_obj = ast.parse(args[0])  # i.e. a dec or exp
 23 |         else:
 24 |             obj = inspect.getsource(args[0])  # i.e. a method
 25 |             ast_obj = ast.parse(obj)
 26 |         json_parsed = method(ast_obj, **kwargs)
 27 |         parsed = json.loads(json_parsed)
 28 | 
 29 |         return parsed
 30 | 
 31 |     return wrapper
 32 | 
 33 | 
 34 | @ast_parse
 35 | def json_ast(node):
 36 |     """Parse an AST object into JSON."""
 37 | 
 38 |     def _format(_node):
 39 |         if isinstance(_node, AST):
 40 |             fields = [("_PyType", _format(_node.__class__.__name__))]
 41 |             fields += [(a, _format(b)) for a, b in iter_fields(_node)]
 42 |             return "{ %s }" % ", ".join(('"%s": %s' % field for field in fields))
 43 |         if isinstance(_node, list):
 44 |             return "[ %s ]" % ", ".join([_format(x) for x in _node])
 45 |         if isinstance(_node, bytes):
 46 |             return json.dumps(_node.decode("utf-8"))
 47 | 
 48 |         return json.dumps(_node)
 49 | 
 50 |     return _format(node)
 51 | 
 52 | 
 53 | def iter_fields(node):
 54 |     """Get attributes of a node."""
 55 |     try:
 56 |         for field in node._fields:
 57 |             yield field, getattr(node, field)
 58 |     except AttributeError:
 59 |         yield
 60 | 
 61 | 
 62 | # ~~~~~~~~~~~~~~~~~~~~~~~~ DRAWING AST ~~~~~~~~~~~~~~~~~~~~~~~~~
 63 | 
 64 | 
 65 | def grapher(graph, ast_nodes, parent_node="", node_hash="__init__"):
 66 |     """Recursively parse JSON-AST object into a tree."""
 67 |     if isinstance(ast_nodes, dict):
 68 |         for key, node in ast_nodes.items():
 69 |             if not parent_node:
 70 |                 parent_node = node
 71 |                 continue
 72 |             if key == "_PyType":
 73 |                 node = graph_detail(node, ast_nodes)  # get node detail for graph
 74 |                 node_hash = draw(parent_node, node, graph=graph, parent_hash=node_hash)
 75 |                 parent_node = node  # once a child now parent
 76 |                 continue
 77 |             # parse recursively
 78 |             if isinstance(node, dict):
 79 |                 grapher(graph, node, parent_node=parent_node, node_hash=node_hash)
 80 |             if isinstance(node, list):
 81 |                 [
 82 |                     grapher(graph, item, parent_node=parent_node, node_hash=node_hash)
 83 |                     for item in node
 84 |                 ]
 85 | 
 86 | 
 87 | def graph_detail(value, ast_scope):
 88 |     """Retrieve node details."""
 89 |     detail_keys = ("module", "n", "s", "id", "name", "attr", "arg")
 90 |     for key in detail_keys:
 91 |         if not isinstance(dict.get(ast_scope, key), type(None)):
 92 |             value = f"{value}\n{key}: {ast_scope[key]}"
 93 | 
 94 |     return value
 95 | 
 96 | 
 97 | def clean_node(method):
 98 |     """Decorator to eliminate illegal characters, check type, and\n
 99 |     shorten lengthy child and parent nodes."""
100 | 
101 |     def wrapper(*args, **kwargs):
102 |         parent_name, child_name = tuple(
103 |             "_node" if node == "node" else node for node in args
104 |         )
105 |         illegal_char = re.compile(r"[,\\/]$")
106 |         illegal_char.sub("*", child_name)
107 |         if not child_name:
108 |             return
109 |         if len(child_name) > 2500:
110 |             child_name = "~~~DOCS: too long to fit on graph~~~"
111 |         args = (parent_name, child_name)
112 | 
113 |         return method(*args, **kwargs)
114 | 
115 |     return wrapper
116 | 
117 | 
118 | @clean_node
119 | def draw(parent_name, child_name, graph, parent_hash):
120 |     """Draw parent and child nodes. Create and return new hash\n
121 |     key declared to a child node."""
122 |     parent_node = pydot.Node(parent_hash, label=parent_name, shape="box")
123 |     child_hash = str(uuid.uuid4())  # create hash key
124 |     child_node = pydot.Node(child_hash, label=child_name, shape="box")
125 | 
126 |     graph.add_node(parent_node)
127 |     graph.add_node(child_node)
128 |     graph.add_edge(pydot.Edge(parent_node, child_node))
129 | 
130 |     return child_hash
131 | 
132 | 
133 | # For jupyter notebooks
134 | def view_tree(pdot):
135 |     """Display tree onto console."""
136 |     tree = Image(pdot.create_png())
137 |     display(tree)
138 | 
139 | 
140 | def parse_input(_input):
141 |     """Parse user input and return an AST-compatible object."""
142 |     try:
143 |         if "." in _input:
144 |             mod, met = _input.split(".")  # handle modules and methods
145 |             module = importlib.import_module(mod)
146 |             method = getattr(module, met)
147 |         else:
148 |             module = importlib.import_module(_input)  # handle modules
149 |             method = module
150 |     except ModuleNotFoundError:
151 |         method = _input  # handle dec, exp
152 | 
153 |     return method
154 | 
155 | 
156 | def main():
157 |     """Take user input and draw an AST.\n
158 |     Save file as PNG."""
159 |     graph = pydot.Dot(
160 |         graph_type="digraph",
161 |         strict=True,
162 |         constraint=True,
163 |         concentrate=True,
164 |         splines="polyline",
165 |     )
166 |     user_input = input("Input a method name, expression, etc.:\n")
167 |     parsed_input = parse_input(user_input)
168 | 
169 |     grapher(graph, json_ast(parsed_input))
170 |     # view_tree(graph)
171 |     if graph.write_png("astree.png"):
172 |         print("Graph made successfully")
173 | 
174 | 
175 | if __name__ == "__main__":
176 |     main()
177 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | appnope==0.1.0
 2 | backcall==0.1.0
 3 | decorator==4.4.2
 4 | ipython==8.10.0
 5 | ipython-genutils==0.2.0
 6 | jedi==0.16.0
 7 | parso==0.6.2
 8 | pexpect==4.8.0
 9 | pickleshare==0.7.5
10 | prompt-toolkit==3.0.5
11 | ptyprocess==0.6.0
12 | pydot-ng==2.0.0
13 | Pygments==2.15.0
14 | pyparsing==2.4.7
15 | six==1.14.0
16 | traitlets==4.3.3
17 | wcwidth==0.1.9
18 | 


--------------------------------------------------------------------------------