| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111 |
- # -*- coding: utf-8 -*-
- """
- 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):
- # Avoid infinite recursion for silly users
- 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()
|