| | |
| | """ |
| | Part of the astor library for Python AST manipulation. |
| | |
| | License: 3-clause BSD |
| | |
| | Copyright (c) 2012-2015 Patrick Maupin |
| | Copyright (c) 2013-2015 Berker Peksag |
| | |
| | Functions that interact with the filesystem go here. |
| | |
| | """ |
| |
|
| | import ast |
| | import sys |
| | import os |
| |
|
| | try: |
| | from tokenize import open as fopen |
| | except ImportError: |
| | fopen = open |
| |
|
| |
|
| | class CodeToAst(object): |
| | """Given a module, or a function that was compiled as part |
| | of a module, re-compile the module into an AST and extract |
| | the sub-AST for the function. Allow caching to reduce |
| | number of compiles. |
| | |
| | Also contains static helper utility functions to |
| | look for python files, to parse python files, and to extract |
| | the file/line information from a code object. |
| | """ |
| |
|
| | @staticmethod |
| | def find_py_files(srctree, ignore=None): |
| | """Return all the python files in a source tree |
| | |
| | Ignores any path that contains the ignore string |
| | |
| | This is not used by other class methods, but is |
| | designed to be used in code that uses this class. |
| | """ |
| |
|
| | if not os.path.isdir(srctree): |
| | yield os.path.split(srctree) |
| | for srcpath, _, fnames in os.walk(srctree): |
| | |
| | if ignore is not None and ignore in srcpath: |
| | continue |
| | for fname in (x for x in fnames if x.endswith('.py')): |
| | yield srcpath, fname |
| |
|
| | @staticmethod |
| | def parse_file(fname): |
| | """Parse a python file into an AST. |
| | |
| | This is a very thin wrapper around ast.parse |
| | |
| | TODO: Handle encodings other than the default for Python 2 |
| | (issue #26) |
| | """ |
| | try: |
| | with fopen(fname) as f: |
| | fstr = f.read() |
| | except IOError: |
| | if fname != 'stdin': |
| | raise |
| | sys.stdout.write('\nReading from stdin:\n\n') |
| | fstr = sys.stdin.read() |
| | fstr = fstr.replace('\r\n', '\n').replace('\r', '\n') |
| | if not fstr.endswith('\n'): |
| | fstr += '\n' |
| | return ast.parse(fstr, filename=fname) |
| |
|
| | @staticmethod |
| | def get_file_info(codeobj): |
| | """Returns the file and line number of a code object. |
| | |
| | If the code object has a __file__ attribute (e.g. if |
| | it is a module), then the returned line number will |
| | be 0 |
| | """ |
| | fname = getattr(codeobj, '__file__', None) |
| | linenum = 0 |
| | if fname is None: |
| | func_code = codeobj.__code__ |
| | fname = func_code.co_filename |
| | linenum = func_code.co_firstlineno |
| | fname = fname.replace('.pyc', '.py') |
| | return fname, linenum |
| |
|
| | def __init__(self, cache=None): |
| | self.cache = cache or {} |
| |
|
| | def __call__(self, codeobj): |
| | cache = self.cache |
| | key = self.get_file_info(codeobj) |
| | result = cache.get(key) |
| | if result is not None: |
| | return result |
| | fname = key[0] |
| | cache[(fname, 0)] = mod_ast = self.parse_file(fname) |
| | for obj in mod_ast.body: |
| | if not isinstance(obj, ast.FunctionDef): |
| | continue |
| | cache[(fname, obj.lineno)] = obj |
| | return cache[key] |
| |
|
| |
|
| | code_to_ast = CodeToAst() |
| |
|
