Skip to content

Instantly share code, notes, and snippets.

@brentp
Forked from fahhem/gist:1036576
Created June 21, 2011 00:31
Show Gist options
  • Save brentp/1036967 to your computer and use it in GitHub Desktop.
Save brentp/1036967 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.functions = kwargs.pop('functions',())
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 _get_function_source(self, fn):
if isinstance(fn, basestring):
fn = globals()[fn]
src = inspect.getsource(fn)
return src
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 = self._get_function_source(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:
if hasattr(mod, "__name__"): mod = mod.__name__
print >> tmp, 'import %s' % mod
for other_fn in self.functions:
print >> tmp, self._get_function_source(other_fn)
# 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'})
def other(a, b):
return a * b
@shedskin((4.5, 12.2), functions=(other,))
def call_other(a, b):
return other(a, b)
print call_other(11.3, 223.2)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment