Inline.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343
  1. from __future__ import absolute_import
  2. import sys, os, re, inspect
  3. import imp
  4. try:
  5. import hashlib
  6. except ImportError:
  7. import md5 as hashlib
  8. from distutils.core import Distribution, Extension
  9. from distutils.command.build_ext import build_ext
  10. import Cython
  11. from ..Compiler.Main import Context, CompilationOptions, default_options
  12. from ..Compiler.ParseTreeTransforms import (CythonTransform,
  13. SkipDeclarations, AnalyseDeclarationsTransform, EnvTransform)
  14. from ..Compiler.TreeFragment import parse_from_strings
  15. from ..Compiler.StringEncoding import _unicode
  16. from .Dependencies import strip_string_literals, cythonize, cached_function
  17. from ..Compiler import Pipeline, Nodes
  18. from ..Utils import get_cython_cache_dir
  19. import cython as cython_module
  20. IS_PY3 = sys.version_info >= (3, 0)
  21. # A utility function to convert user-supplied ASCII strings to unicode.
  22. if sys.version_info[0] < 3:
  23. def to_unicode(s):
  24. if isinstance(s, bytes):
  25. return s.decode('ascii')
  26. else:
  27. return s
  28. else:
  29. to_unicode = lambda x: x
  30. class UnboundSymbols(EnvTransform, SkipDeclarations):
  31. def __init__(self):
  32. CythonTransform.__init__(self, None)
  33. self.unbound = set()
  34. def visit_NameNode(self, node):
  35. if not self.current_env().lookup(node.name):
  36. self.unbound.add(node.name)
  37. return node
  38. def __call__(self, node):
  39. super(UnboundSymbols, self).__call__(node)
  40. return self.unbound
  41. @cached_function
  42. def unbound_symbols(code, context=None):
  43. code = to_unicode(code)
  44. if context is None:
  45. context = Context([], default_options)
  46. from ..Compiler.ParseTreeTransforms import AnalyseDeclarationsTransform
  47. tree = parse_from_strings('(tree fragment)', code)
  48. for phase in Pipeline.create_pipeline(context, 'pyx'):
  49. if phase is None:
  50. continue
  51. tree = phase(tree)
  52. if isinstance(phase, AnalyseDeclarationsTransform):
  53. break
  54. try:
  55. import builtins
  56. except ImportError:
  57. import __builtin__ as builtins
  58. return tuple(UnboundSymbols()(tree) - set(dir(builtins)))
  59. def unsafe_type(arg, context=None):
  60. py_type = type(arg)
  61. if py_type is int:
  62. return 'long'
  63. else:
  64. return safe_type(arg, context)
  65. def safe_type(arg, context=None):
  66. py_type = type(arg)
  67. if py_type in (list, tuple, dict, str):
  68. return py_type.__name__
  69. elif py_type is complex:
  70. return 'double complex'
  71. elif py_type is float:
  72. return 'double'
  73. elif py_type is bool:
  74. return 'bint'
  75. elif 'numpy' in sys.modules and isinstance(arg, sys.modules['numpy'].ndarray):
  76. return 'numpy.ndarray[numpy.%s_t, ndim=%s]' % (arg.dtype.name, arg.ndim)
  77. else:
  78. for base_type in py_type.mro():
  79. if base_type.__module__ in ('__builtin__', 'builtins'):
  80. return 'object'
  81. module = context.find_module(base_type.__module__, need_pxd=False)
  82. if module:
  83. entry = module.lookup(base_type.__name__)
  84. if entry.is_type:
  85. return '%s.%s' % (base_type.__module__, base_type.__name__)
  86. return 'object'
  87. def _get_build_extension():
  88. dist = Distribution()
  89. # Ensure the build respects distutils configuration by parsing
  90. # the configuration files
  91. config_files = dist.find_config_files()
  92. dist.parse_config_files(config_files)
  93. build_extension = build_ext(dist)
  94. build_extension.finalize_options()
  95. return build_extension
  96. @cached_function
  97. def _create_context(cython_include_dirs):
  98. return Context(list(cython_include_dirs), default_options)
  99. _cython_inline_cache = {}
  100. _cython_inline_default_context = _create_context(('.',))
  101. def _populate_unbound(kwds, unbound_symbols, locals=None, globals=None):
  102. for symbol in unbound_symbols:
  103. if symbol not in kwds:
  104. if locals is None or globals is None:
  105. calling_frame = inspect.currentframe().f_back.f_back.f_back
  106. if locals is None:
  107. locals = calling_frame.f_locals
  108. if globals is None:
  109. globals = calling_frame.f_globals
  110. if symbol in locals:
  111. kwds[symbol] = locals[symbol]
  112. elif symbol in globals:
  113. kwds[symbol] = globals[symbol]
  114. else:
  115. print("Couldn't find %r" % symbol)
  116. def cython_inline(code, get_type=unsafe_type, lib_dir=os.path.join(get_cython_cache_dir(), 'inline'),
  117. cython_include_dirs=None, force=False, quiet=False, locals=None, globals=None, **kwds):
  118. if get_type is None:
  119. get_type = lambda x: 'object'
  120. ctx = _create_context(tuple(cython_include_dirs)) if cython_include_dirs else _cython_inline_default_context
  121. # Fast path if this has been called in this session.
  122. _unbound_symbols = _cython_inline_cache.get(code)
  123. if _unbound_symbols is not None:
  124. _populate_unbound(kwds, _unbound_symbols, locals, globals)
  125. args = sorted(kwds.items())
  126. arg_sigs = tuple([(get_type(value, ctx), arg) for arg, value in args])
  127. invoke = _cython_inline_cache.get((code, arg_sigs))
  128. if invoke is not None:
  129. arg_list = [arg[1] for arg in args]
  130. return invoke(*arg_list)
  131. orig_code = code
  132. code = to_unicode(code)
  133. code, literals = strip_string_literals(code)
  134. code = strip_common_indent(code)
  135. if locals is None:
  136. locals = inspect.currentframe().f_back.f_back.f_locals
  137. if globals is None:
  138. globals = inspect.currentframe().f_back.f_back.f_globals
  139. try:
  140. _cython_inline_cache[orig_code] = _unbound_symbols = unbound_symbols(code)
  141. _populate_unbound(kwds, _unbound_symbols, locals, globals)
  142. except AssertionError:
  143. if not quiet:
  144. # Parsing from strings not fully supported (e.g. cimports).
  145. print("Could not parse code as a string (to extract unbound symbols).")
  146. cimports = []
  147. for name, arg in list(kwds.items()):
  148. if arg is cython_module:
  149. cimports.append('\ncimport cython as %s' % name)
  150. del kwds[name]
  151. arg_names = sorted(kwds)
  152. arg_sigs = tuple([(get_type(kwds[arg], ctx), arg) for arg in arg_names])
  153. key = orig_code, arg_sigs, sys.version_info, sys.executable, Cython.__version__
  154. module_name = "_cython_inline_" + hashlib.md5(_unicode(key).encode('utf-8')).hexdigest()
  155. if module_name in sys.modules:
  156. module = sys.modules[module_name]
  157. else:
  158. build_extension = None
  159. if cython_inline.so_ext is None:
  160. # Figure out and cache current extension suffix
  161. build_extension = _get_build_extension()
  162. cython_inline.so_ext = build_extension.get_ext_filename('')
  163. module_path = os.path.join(lib_dir, module_name + cython_inline.so_ext)
  164. if not os.path.exists(lib_dir):
  165. os.makedirs(lib_dir)
  166. if force or not os.path.isfile(module_path):
  167. cflags = []
  168. c_include_dirs = []
  169. qualified = re.compile(r'([.\w]+)[.]')
  170. for type, _ in arg_sigs:
  171. m = qualified.match(type)
  172. if m:
  173. cimports.append('\ncimport %s' % m.groups()[0])
  174. # one special case
  175. if m.groups()[0] == 'numpy':
  176. import numpy
  177. c_include_dirs.append(numpy.get_include())
  178. # cflags.append('-Wno-unused')
  179. module_body, func_body = extract_func_code(code)
  180. params = ', '.join(['%s %s' % a for a in arg_sigs])
  181. module_code = """
  182. %(module_body)s
  183. %(cimports)s
  184. def __invoke(%(params)s):
  185. %(func_body)s
  186. return locals()
  187. """ % {'cimports': '\n'.join(cimports),
  188. 'module_body': module_body,
  189. 'params': params,
  190. 'func_body': func_body }
  191. for key, value in literals.items():
  192. module_code = module_code.replace(key, value)
  193. pyx_file = os.path.join(lib_dir, module_name + '.pyx')
  194. fh = open(pyx_file, 'w')
  195. try:
  196. fh.write(module_code)
  197. finally:
  198. fh.close()
  199. extension = Extension(
  200. name = module_name,
  201. sources = [pyx_file],
  202. include_dirs = c_include_dirs,
  203. extra_compile_args = cflags)
  204. if build_extension is None:
  205. build_extension = _get_build_extension()
  206. build_extension.extensions = cythonize([extension], include_path=cython_include_dirs or ['.'], quiet=quiet)
  207. build_extension.build_temp = os.path.dirname(pyx_file)
  208. build_extension.build_lib = lib_dir
  209. build_extension.run()
  210. module = imp.load_dynamic(module_name, module_path)
  211. _cython_inline_cache[orig_code, arg_sigs] = module.__invoke
  212. arg_list = [kwds[arg] for arg in arg_names]
  213. return module.__invoke(*arg_list)
  214. # Cached suffix used by cython_inline above. None should get
  215. # overridden with actual value upon the first cython_inline invocation
  216. cython_inline.so_ext = None
  217. _find_non_space = re.compile('[^ ]').search
  218. def strip_common_indent(code):
  219. min_indent = None
  220. lines = code.splitlines()
  221. for line in lines:
  222. match = _find_non_space(line)
  223. if not match:
  224. continue # blank
  225. indent = match.start()
  226. if line[indent] == '#':
  227. continue # comment
  228. if min_indent is None or min_indent > indent:
  229. min_indent = indent
  230. for ix, line in enumerate(lines):
  231. match = _find_non_space(line)
  232. if not match or not line or line[indent:indent+1] == '#':
  233. continue
  234. lines[ix] = line[min_indent:]
  235. return '\n'.join(lines)
  236. module_statement = re.compile(r'^((cdef +(extern|class))|cimport|(from .+ cimport)|(from .+ import +[*]))')
  237. def extract_func_code(code):
  238. module = []
  239. function = []
  240. current = function
  241. code = code.replace('\t', ' ')
  242. lines = code.split('\n')
  243. for line in lines:
  244. if not line.startswith(' '):
  245. if module_statement.match(line):
  246. current = module
  247. else:
  248. current = function
  249. current.append(line)
  250. return '\n'.join(module), ' ' + '\n '.join(function)
  251. try:
  252. from inspect import getcallargs
  253. except ImportError:
  254. def getcallargs(func, *arg_values, **kwd_values):
  255. all = {}
  256. args, varargs, kwds, defaults = inspect.getargspec(func)
  257. if varargs is not None:
  258. all[varargs] = arg_values[len(args):]
  259. for name, value in zip(args, arg_values):
  260. all[name] = value
  261. for name, value in list(kwd_values.items()):
  262. if name in args:
  263. if name in all:
  264. raise TypeError("Duplicate argument %s" % name)
  265. all[name] = kwd_values.pop(name)
  266. if kwds is not None:
  267. all[kwds] = kwd_values
  268. elif kwd_values:
  269. raise TypeError("Unexpected keyword arguments: %s" % list(kwd_values))
  270. if defaults is None:
  271. defaults = ()
  272. first_default = len(args) - len(defaults)
  273. for ix, name in enumerate(args):
  274. if name not in all:
  275. if ix >= first_default:
  276. all[name] = defaults[ix - first_default]
  277. else:
  278. raise TypeError("Missing argument: %s" % name)
  279. return all
  280. def get_body(source):
  281. ix = source.index(':')
  282. if source[:5] == 'lambda':
  283. return "return %s" % source[ix+1:]
  284. else:
  285. return source[ix+1:]
  286. # Lots to be done here... It would be especially cool if compiled functions
  287. # could invoke each other quickly.
  288. class RuntimeCompiledFunction(object):
  289. def __init__(self, f):
  290. self._f = f
  291. self._body = get_body(inspect.getsource(f))
  292. def __call__(self, *args, **kwds):
  293. all = getcallargs(self._f, *args, **kwds)
  294. if IS_PY3:
  295. return cython_inline(self._body, locals=self._f.__globals__, globals=self._f.__globals__, **all)
  296. else:
  297. return cython_inline(self._body, locals=self._f.func_globals, globals=self._f.func_globals, **all)