file_util.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. # -*- coding: utf-8 -*-
  2. """
  3. Part of the astor library for Python AST manipulation.
  4. License: 3-clause BSD
  5. Copyright (c) 2012-2015 Patrick Maupin
  6. Copyright (c) 2013-2015 Berker Peksag
  7. Functions that interact with the filesystem go here.
  8. """
  9. import ast
  10. import sys
  11. import os
  12. try:
  13. from tokenize import open as fopen
  14. except ImportError:
  15. fopen = open
  16. class CodeToAst(object):
  17. """Given a module, or a function that was compiled as part
  18. of a module, re-compile the module into an AST and extract
  19. the sub-AST for the function. Allow caching to reduce
  20. number of compiles.
  21. Also contains static helper utility functions to
  22. look for python files, to parse python files, and to extract
  23. the file/line information from a code object.
  24. """
  25. @staticmethod
  26. def find_py_files(srctree, ignore=None):
  27. """Return all the python files in a source tree
  28. Ignores any path that contains the ignore string
  29. This is not used by other class methods, but is
  30. designed to be used in code that uses this class.
  31. """
  32. if not os.path.isdir(srctree):
  33. yield os.path.split(srctree)
  34. for srcpath, _, fnames in os.walk(srctree):
  35. # Avoid infinite recursion for silly users
  36. if ignore is not None and ignore in srcpath:
  37. continue
  38. for fname in (x for x in fnames if x.endswith('.py')):
  39. yield srcpath, fname
  40. @staticmethod
  41. def parse_file(fname):
  42. """Parse a python file into an AST.
  43. This is a very thin wrapper around ast.parse
  44. TODO: Handle encodings other than the default for Python 2
  45. (issue #26)
  46. """
  47. try:
  48. with fopen(fname) as f:
  49. fstr = f.read()
  50. except IOError:
  51. if fname != 'stdin':
  52. raise
  53. sys.stdout.write('\nReading from stdin:\n\n')
  54. fstr = sys.stdin.read()
  55. fstr = fstr.replace('\r\n', '\n').replace('\r', '\n')
  56. if not fstr.endswith('\n'):
  57. fstr += '\n'
  58. return ast.parse(fstr, filename=fname)
  59. @staticmethod
  60. def get_file_info(codeobj):
  61. """Returns the file and line number of a code object.
  62. If the code object has a __file__ attribute (e.g. if
  63. it is a module), then the returned line number will
  64. be 0
  65. """
  66. fname = getattr(codeobj, '__file__', None)
  67. linenum = 0
  68. if fname is None:
  69. func_code = codeobj.__code__
  70. fname = func_code.co_filename
  71. linenum = func_code.co_firstlineno
  72. fname = fname.replace('.pyc', '.py')
  73. return fname, linenum
  74. def __init__(self, cache=None):
  75. self.cache = cache or {}
  76. def __call__(self, codeobj):
  77. cache = self.cache
  78. key = self.get_file_info(codeobj)
  79. result = cache.get(key)
  80. if result is not None:
  81. return result
  82. fname = key[0]
  83. cache[(fname, 0)] = mod_ast = self.parse_file(fname)
  84. for obj in mod_ast.body:
  85. if not isinstance(obj, ast.FunctionDef):
  86. continue
  87. cache[(fname, obj.lineno)] = obj
  88. return cache[key]
  89. code_to_ast = CodeToAst()