Skip to content

Instantly share code, notes, and snippets.

@fahhem
Forked from brentp/gist:1036011
Created June 20, 2011 21:03
Show Gist options
  • Save fahhem/1036576 to your computer and use it in GitHub Desktop.
Save fahhem/1036576 to your computer and use it in GitHub Desktop.
shedskin decorator
import shedskin
import os.path as op
import sys
from shedskin import setgx, newgx, getgx, infer, annotate, cpp, shared
import inspect
import hashlib
import subprocess
class shedskin(object):
r"""
decorator to shedskin-ify functions.
and example is::
@shedskin((2, 4), (6, 12), long=True, random=True, wrap_around_check=False)
def adder(a, b):
return a + b
print adder(6, 8)
note the decorator is called with example invocations so that
shedskin can infer types. After the first invocation, the extension
module is stored so that subsequent calls will run quickly.
If either the function or the decorator is changed, a new module
will be recreated.
"""
kwlookup = {
'nobounds': 'bounds_checking',
'long': 'longlong',
'nowrap': 'wrap_around_check',
'random': 'fast_random',
'strhash': 'fast_hash',
}
def __init__(self, *invocations, **kwargs):
self.modules = kwargs.pop('modules',())
self.kwargs = kwargs
self.invocations = invocations
def _hash(self, source):
return hashlib.md5(source).hexdigest()
def _tmp(self, source_hash, ext=".py"):
return "shedskin" + source_hash + ext
def __call__(self, fn):
setgx(newgx())
# set kwargs from the __init__ call.
for k, v in self.kwargs.items():
k = shedskin.kwlookup.get(k, k)
setattr(getgx(), k, v)
getgx().annotation = True
getgx().extension_module = True
src = inspect.getsource(fn)
source_hash = self._hash(src)
if self._is_up_to_date(source_hash):
mod = self._get_module(self._tmp(source_hash))
return getattr(mod, fn.func_name)
tmp = open(self._tmp(source_hash), "w")
for mod in self.modules:
print >> tmp, 'import %s' % mod
# hack to get the function source without the decorator line...
# needs to be fixed...
print >> tmp, src.split('\n',1)[1]
for i in self.invocations:
print >>tmp, "%s%s" % (fn.func_name, str(i))
tmp.close()
makefile = getgx().makefile_name = "Makefile_%s" % source_hash
self._run_shedskin(tmp.name, makefile)
mod = self._get_module(tmp.name)
return getattr(mod, fn.func_name)
def _is_up_to_date(self, source_hash):
return op.exists(self._tmp(source_hash, ext=".so"))
def _run_shedskin(self, name, makefile):
old = sys.stdout
log = sys.stdout = open(name + ".log", "w")
getgx().main_mod = name[:-3]
infer.analyze(name)
annotate.annotate()
cpp.generate_code()
shared.print_errors()
ret = subprocess.call("make -f %s" % makefile, shell=True,
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
log.close()
sys.stdout = old
if ret != 0:
print >>sys.stderr, "error making %s" % makefile
print open(log.name).read()
def _get_module(self, name):
if name.endswith(".py"):
name = name[:-3]
mod = __import__(name)
return mod
@shedskin((6, 8), long=True, random=True, wrap_around_check=False)
def adder(a, b):
return a + b
print adder(11, 8)
import collections
@shedskin(({'a':'b'},),modules=('collections',))
def makedefdict(a):
dd = collections.defaultdict(str)
for b in a:
dd[b] = a[b]
return dd
print makedefdict({'c':'d'})
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment