self_scope.py

""" Example:

In [5]: class Class:
   ...:     def __init__(self, a, b):
   ...:         self.a = a
   ...:         self.b = b
   ...: 
   ...:     @self_scope
   ...:     def f(self):
   ...:         return a + b
   ...: 
   ...:     @self_scope
   ...:     def g(self):
   ...:         return a + c
   ...:     

In [6]: inst = Class(1, 2)

In [7]: inst.f()
Out[7]: 3

In [8]: inst.g()
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-166ea1fc502a> in _lookup(attr, ob)
      5         try:
----> 6             return getattr(ob, attr)
      7         except AttributeError:

AttributeError: 'Class' object has no attribute 'c'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-4-166ea1fc502a> in _lookup(attr, ob)
      8             try:
----> 9                 sys._getframe(1).f_globals[attr]
     10             except KeyError:

KeyError: 'c'

During handling of the above exception, another exception occurred:

NameError                                 Traceback (most recent call last)
<ipython-input-8-79dc91e4b37c> in <module>()
----> 1 inst.g()

<ipython-input-5-5da52f677f4e> in g(self)
     10     @self_scope
     11     def g(self):
---> 12         return a + c
     13 

<ipython-input-4-166ea1fc502a> in _lookup(attr, ob)
      9                 sys._getframe(1).f_globals[attr]
     10             except KeyError:
---> 11                 raise NameError(attr)
     12 
     13     @pattern(instrs.LOAD_GLOBAL)

NameError: c
"""
from functools import partial
import sys

from codetransformer import instructions as instrs, CodeTransformer, pattern


@type.__call__
class self_scope(CodeTransformer):
    @staticmethod
    def _lookup(attr, ob):
        try:
            return getattr(ob, attr)
        except AttributeError:
            try:
                sys._getframe(1).f_globals[attr]
            except KeyError:
                raise NameError(attr)
            
    @pattern(instrs.LOAD_GLOBAL)
    def _(self, instr):
        yield instrs.LOAD_CONST(partial(self._lookup, instr.arg)).steal(instr)
        yield instrs.LOAD_FAST('self')
        yield instrs.CALL_FUNCTION(1)