cmlenz@336: # -*- coding: utf-8 -*-
cmlenz@336: #
cmlenz@408: # Copyright (C) 2006-2007 Edgewall Software
cmlenz@336: # All rights reserved.
cmlenz@336: #
cmlenz@336: # This software is licensed as described in the file COPYING, which
cmlenz@336: # you should have received as part of this distribution. The terms
cmlenz@336: # are also available at http://genshi.edgewall.org/wiki/License.
cmlenz@336: #
cmlenz@336: # This software consists of voluntary contributions made by many
cmlenz@336: # individuals. For the exact contribution history, see the revision
cmlenz@336: # history and logs, available at http://genshi.edgewall.org/log/.
cmlenz@336: 
cmlenz@336: """Support for "safe" evaluation of Python expressions."""
cmlenz@336: 
cmlenz@336: import __builtin__
cmlenz@336: from compiler import ast, parse
cmlenz@405: from compiler.pycodegen import ExpressionCodeGenerator, ModuleCodeGenerator
cmlenz@336: import new
cmlenz@357: try:
cmlenz@357:     set
cmlenz@357: except NameError:
cmlenz@565:     from sets import ImmutableSet as frozenset
cmlenz@357:     from sets import Set as set
cmlenz@405: import sys
cmlenz@601: from textwrap import dedent
aronacher@642: from types import FunctionType, MethodType
cmlenz@357: 
cmlenz@401: from genshi.core import Markup
cmlenz@418: from genshi.template.base import TemplateRuntimeError
aronacher@643: from genshi.util import flatten, safe_range
cmlenz@336: 
cmlenz@442: __all__ = ['Code', 'Expression', 'Suite', 'LenientLookup', 'StrictLookup',
cmlenz@442:            'Undefined', 'UndefinedError']
cmlenz@425: __docformat__ = 'restructuredtext en'
cmlenz@336: 
cmlenz@336: 
cmlenz@405: class Code(object):
cmlenz@405:     """Abstract base class for the `Expression` and `Suite` classes."""
aronacher@642:     __slots__ = ['source', 'code', 'ast', 'secure', '_globals']
cmlenz@405: 
cmlenz@606:     def __init__(self, source, filename=None, lineno=-1, lookup='strict',
aronacher@642:                  xform=None, secure=False):
cmlenz@405:         """Create the code object, either from a string, or from an AST node.
cmlenz@405:         
cmlenz@425:         :param source: either a string containing the source code, or an AST
cmlenz@425:                        node
cmlenz@425:         :param filename: the (preferably absolute) name of the file containing
cmlenz@425:                          the code
cmlenz@425:         :param lineno: the number of the line on which the code was found
cmlenz@442:         :param lookup: the lookup class that defines how variables are looked
cmlenz@606:                        up in the context; can be either "strict" (the default),
cmlenz@606:                        "lenient", or a custom lookup class
cmlenz@604:         :param xform: the AST transformer that should be applied to the code;
cmlenz@604:                       if `None`, the appropriate transformation is chosen
cmlenz@604:                       depending on the mode
aronacher@642:         :param secure: If security features should be enabled.
cmlenz@405:         """
cmlenz@405:         if isinstance(source, basestring):
cmlenz@405:             self.source = source
cmlenz@405:             node = _parse(source, mode=self.mode)
cmlenz@405:         else:
cmlenz@601:             assert isinstance(source, ast.Node), \
cmlenz@601:                 'Expected string or AST node, but got %r' % source
cmlenz@405:             self.source = '?'
cmlenz@405:             if self.mode == 'eval':
cmlenz@405:                 node = ast.Expression(source)
cmlenz@405:             else:
cmlenz@405:                 node = ast.Module(None, source)
cmlenz@405: 
cmlenz@565:         self.ast = node
cmlenz@405:         self.code = _compile(node, self.source, mode=self.mode,
aronacher@642:                              filename=filename, lineno=lineno, xform=xform,
aronacher@642:                              secure=secure)
cmlenz@442:         if lookup is None:
cmlenz@442:             lookup = LenientLookup
cmlenz@442:         elif isinstance(lookup, basestring):
aronacher@642:             lookup = {
aronacher@642:                 'lenient':  LenientLookup,
aronacher@642:                 'strict':   StrictLookup
aronacher@642:             }[lookup]
aronacher@642:             if secure:
aronacher@642:                 lookup = SecurityLookupWrapper(lookup)
aronacher@642:         self.secure = secure
cmlenz@442:         self._globals = lookup.globals()
cmlenz@405: 
cmlenz@405:     def __eq__(self, other):
cmlenz@405:         return (type(other) == type(self)) and (self.code == other.code)
cmlenz@405: 
cmlenz@405:     def __hash__(self):
cmlenz@405:         return hash(self.code)
cmlenz@405: 
cmlenz@405:     def __ne__(self, other):
cmlenz@405:         return not self == other
cmlenz@405: 
cmlenz@405:     def __repr__(self):
cmlenz@405:         return '%s(%r)' % (self.__class__.__name__, self.source)
cmlenz@405: 
cmlenz@405: 
cmlenz@405: class Expression(Code):
cmlenz@336:     """Evaluates Python expressions used in templates.
cmlenz@336: 
cmlenz@336:     >>> data = dict(test='Foo', items=[1, 2, 3], dict={'some': 'thing'})
cmlenz@336:     >>> Expression('test').evaluate(data)
cmlenz@336:     'Foo'
cmlenz@336: 
cmlenz@336:     >>> Expression('items[0]').evaluate(data)
cmlenz@336:     1
cmlenz@336:     >>> Expression('items[-1]').evaluate(data)
cmlenz@336:     3
cmlenz@336:     >>> Expression('dict["some"]').evaluate(data)
cmlenz@336:     'thing'
cmlenz@336:     
cmlenz@336:     Similar to e.g. Javascript, expressions in templates can use the dot
cmlenz@336:     notation for attribute access to access items in mappings:
cmlenz@336:     
cmlenz@336:     >>> Expression('dict.some').evaluate(data)
cmlenz@336:     'thing'
cmlenz@336:     
cmlenz@336:     This also works the other way around: item access can be used to access
cmlenz@425:     any object attribute:
cmlenz@336:     
cmlenz@336:     >>> class MyClass(object):
cmlenz@336:     ...     myattr = 'Bar'
cmlenz@336:     >>> data = dict(mine=MyClass(), key='myattr')
cmlenz@336:     >>> Expression('mine.myattr').evaluate(data)
cmlenz@336:     'Bar'
cmlenz@336:     >>> Expression('mine["myattr"]').evaluate(data)
cmlenz@336:     'Bar'
cmlenz@336:     >>> Expression('mine[key]').evaluate(data)
cmlenz@336:     'Bar'
cmlenz@336:     
cmlenz@336:     All of the standard Python operators are available to template expressions.
cmlenz@425:     Built-in functions such as ``len()`` are also available in template
cmlenz@336:     expressions:
cmlenz@336:     
cmlenz@336:     >>> data = dict(items=[1, 2, 3])
cmlenz@336:     >>> Expression('len(items)').evaluate(data)
cmlenz@336:     3
cmlenz@336:     """
cmlenz@405:     __slots__ = []
cmlenz@405:     mode = 'eval'
cmlenz@336: 
cmlenz@343:     def evaluate(self, data):
cmlenz@336:         """Evaluate the expression against the given data dictionary.
cmlenz@336:         
cmlenz@425:         :param data: a mapping containing the data to evaluate against
cmlenz@425:         :return: the result of the evaluation
cmlenz@336:         """
cmlenz@418:         __traceback_hide__ = 'before_and_this'
cmlenz@442:         _globals = self._globals
cmlenz@442:         _globals['data'] = data
cmlenz@442:         return eval(self.code, _globals, {'data': data})
cmlenz@336: 
cmlenz@336: 
cmlenz@405: class Suite(Code):
cmlenz@405:     """Executes Python statements used in templates.
cmlenz@405: 
cmlenz@405:     >>> data = dict(test='Foo', items=[1, 2, 3], dict={'some': 'thing'})
cmlenz@473:     >>> Suite("foo = dict['some']").execute(data)
cmlenz@405:     >>> data['foo']
cmlenz@405:     'thing'
cmlenz@405:     """
cmlenz@405:     __slots__ = []
cmlenz@405:     mode = 'exec'
cmlenz@405: 
cmlenz@405:     def execute(self, data):
cmlenz@405:         """Execute the suite in the given data dictionary.
cmlenz@405:         
cmlenz@425:         :param data: a mapping containing the data to execute in
cmlenz@405:         """
cmlenz@418:         __traceback_hide__ = 'before_and_this'
cmlenz@442:         _globals = self._globals
cmlenz@442:         _globals['data'] = data
cmlenz@442:         exec self.code in _globals, data
cmlenz@405: 
cmlenz@405: 
cmlenz@442: UNDEFINED = object()
cmlenz@442: 
cmlenz@442: 
cmlenz@442: class UndefinedError(TemplateRuntimeError):
cmlenz@442:     """Exception thrown when a template expression attempts to access a variable
cmlenz@442:     not defined in the context.
cmlenz@442:     
cmlenz@442:     :see: `LenientLookup`, `StrictLookup`
cmlenz@442:     """
cmlenz@442:     def __init__(self, name, owner=UNDEFINED):
cmlenz@442:         if owner is not UNDEFINED:
cmlenz@442:             message = '%s has no member named "%s"' % (repr(owner), name)
cmlenz@442:         else:
cmlenz@442:             message = '"%s" not defined' % name
cmlenz@442:         TemplateRuntimeError.__init__(self, message)
cmlenz@442: 
cmlenz@442: 
cmlenz@442: class Undefined(object):
cmlenz@442:     """Represents a reference to an undefined variable.
cmlenz@442:     
cmlenz@442:     Unlike the Python runtime, template expressions can refer to an undefined
cmlenz@442:     variable without causing a `NameError` to be raised. The result will be an
cmlenz@442:     instance of the `Undefined` class, which is treated the same as ``False`` in
cmlenz@442:     conditions, but raise an exception on any other operation:
cmlenz@442:     
cmlenz@442:     >>> foo = Undefined('foo')
cmlenz@442:     >>> bool(foo)
cmlenz@442:     False
cmlenz@442:     >>> list(foo)
cmlenz@442:     []
cmlenz@442:     >>> print foo
cmlenz@442:     undefined
cmlenz@442:     
cmlenz@442:     However, calling an undefined variable, or trying to access an attribute
cmlenz@442:     of that variable, will raise an exception that includes the name used to
cmlenz@442:     reference that undefined variable.
cmlenz@442:     
cmlenz@442:     >>> foo('bar')
cmlenz@442:     Traceback (most recent call last):
cmlenz@442:         ...
cmlenz@442:     UndefinedError: "foo" not defined
cmlenz@442: 
cmlenz@442:     >>> foo.bar
cmlenz@442:     Traceback (most recent call last):
cmlenz@442:         ...
cmlenz@442:     UndefinedError: "foo" not defined
cmlenz@442:     
cmlenz@442:     :see: `LenientLookup`
cmlenz@442:     """
cmlenz@442:     __slots__ = ['_name', '_owner']
cmlenz@442: 
cmlenz@442:     def __init__(self, name, owner=UNDEFINED):
cmlenz@442:         """Initialize the object.
cmlenz@442:         
cmlenz@442:         :param name: the name of the reference
cmlenz@442:         :param owner: the owning object, if the variable is accessed as a member
cmlenz@418:         """
cmlenz@442:         self._name = name
cmlenz@442:         self._owner = owner
cmlenz@336: 
cmlenz@442:     def __iter__(self):
cmlenz@442:         return iter([])
cmlenz@442: 
cmlenz@442:     def __nonzero__(self):
cmlenz@442:         return False
cmlenz@442: 
cmlenz@442:     def __repr__(self):
cmlenz@442:         return '<%s %r>' % (self.__class__.__name__, self._name)
cmlenz@442: 
cmlenz@442:     def __str__(self):
cmlenz@442:         return 'undefined'
cmlenz@442: 
cmlenz@442:     def _die(self, *args, **kwargs):
cmlenz@442:         """Raise an `UndefinedError`."""
cmlenz@442:         __traceback_hide__ = True
cmlenz@442:         raise UndefinedError(self._name, self._owner)
cmlenz@442:     __call__ = __getattr__ = __getitem__ = _die
cmlenz@442: 
cmlenz@442: 
cmlenz@442: class LookupBase(object):
cmlenz@442:     """Abstract base class for variable lookup implementations."""
cmlenz@442: 
cmlenz@442:     def globals(cls):
cmlenz@442:         """Construct the globals dictionary to use as the execution context for
cmlenz@442:         the expression or suite.
cmlenz@418:         """
cmlenz@442:         return {
cmlenz@442:             '_lookup_name': cls.lookup_name,
cmlenz@442:             '_lookup_attr': cls.lookup_attr,
cmlenz@579:             '_lookup_item': cls.lookup_item,
cmlenz@579:             'UndefinedError': UndefinedError
cmlenz@442:         }
cmlenz@442:     globals = classmethod(globals)
cmlenz@442: 
cmlenz@442:     def lookup_name(cls, data, name):
cmlenz@442:         __traceback_hide__ = True
cmlenz@442:         val = data.get(name, UNDEFINED)
cmlenz@442:         if val is UNDEFINED:
cmlenz@442:             val = BUILTINS.get(name, val)
cmlenz@442:             if val is UNDEFINED:
cmlenz@569:                 val = cls.undefined(name)
cmlenz@442:         return val
cmlenz@442:     lookup_name = classmethod(lookup_name)
cmlenz@442: 
cmlenz@442:     def lookup_attr(cls, data, obj, key):
cmlenz@442:         __traceback_hide__ = True
cmlenz@569:         val = getattr(obj, key, UNDEFINED)
cmlenz@569:         if val is UNDEFINED:
cmlenz@569:             try:
cmlenz@569:                 val = obj[key]
cmlenz@569:             except (KeyError, TypeError):
cmlenz@569:                 val = cls.undefined(key, owner=obj)
cmlenz@569:         return val
cmlenz@442:     lookup_attr = classmethod(lookup_attr)
cmlenz@442: 
cmlenz@442:     def lookup_item(cls, data, obj, key):
cmlenz@442:         __traceback_hide__ = True
cmlenz@442:         if len(key) == 1:
cmlenz@442:             key = key[0]
cmlenz@442:         try:
cmlenz@442:             return obj[key]
cmlenz@442:         except (AttributeError, KeyError, IndexError, TypeError), e:
cmlenz@442:             if isinstance(key, basestring):
cmlenz@442:                 val = getattr(obj, key, UNDEFINED)
cmlenz@442:                 if val is UNDEFINED:
cmlenz@569:                     val = cls.undefined(key, owner=obj)
cmlenz@442:                 return val
cmlenz@442:             raise
cmlenz@442:     lookup_item = classmethod(lookup_item)
cmlenz@442: 
cmlenz@442:     def undefined(cls, key, owner=UNDEFINED):
cmlenz@442:         """Can be overridden by subclasses to specify behavior when undefined
cmlenz@442:         variables are accessed.
cmlenz@442:         
cmlenz@442:         :param key: the name of the variable
cmlenz@442:         :param owner: the owning object, if the variable is accessed as a member
cmlenz@442:         """
cmlenz@442:         raise NotImplementedError
cmlenz@442:     undefined = classmethod(undefined)
cmlenz@442: 
cmlenz@442: 
cmlenz@442: class LenientLookup(LookupBase):
cmlenz@442:     """Default variable lookup mechanism for expressions.
cmlenz@442:     
cmlenz@442:     When an undefined variable is referenced using this lookup style, the
cmlenz@442:     reference evaluates to an instance of the `Undefined` class:
cmlenz@442:     
cmlenz@442:     >>> expr = Expression('nothing', lookup='lenient')
cmlenz@442:     >>> undef = expr.evaluate({})
cmlenz@442:     >>> undef
cmlenz@442:     <Undefined 'nothing'>
cmlenz@442:     
cmlenz@442:     The same will happen when a non-existing attribute or item is accessed on
cmlenz@442:     an existing object:
cmlenz@442:     
cmlenz@442:     >>> expr = Expression('something.nil', lookup='lenient')
cmlenz@442:     >>> expr.evaluate({'something': dict()})
cmlenz@442:     <Undefined 'nil'>
cmlenz@442:     
cmlenz@442:     See the documentation of the `Undefined` class for details on the behavior
cmlenz@442:     of such objects.
cmlenz@442:     
cmlenz@442:     :see: `StrictLookup`
cmlenz@442:     """
cmlenz@442:     def undefined(cls, key, owner=UNDEFINED):
cmlenz@442:         """Return an ``Undefined`` object."""
cmlenz@442:         __traceback_hide__ = True
cmlenz@442:         return Undefined(key, owner=owner)
cmlenz@442:     undefined = classmethod(undefined)
cmlenz@442: 
cmlenz@442: 
cmlenz@442: class StrictLookup(LookupBase):
cmlenz@442:     """Strict variable lookup mechanism for expressions.
cmlenz@442:     
cmlenz@442:     Referencing an undefined variable using this lookup style will immediately
cmlenz@442:     raise an ``UndefinedError``:
cmlenz@442:     
cmlenz@442:     >>> expr = Expression('nothing', lookup='strict')
cmlenz@442:     >>> expr.evaluate({})
cmlenz@442:     Traceback (most recent call last):
cmlenz@442:         ...
cmlenz@442:     UndefinedError: "nothing" not defined
cmlenz@442:     
cmlenz@442:     The same happens when a non-existing attribute or item is accessed on an
cmlenz@442:     existing object:
cmlenz@442:     
cmlenz@442:     >>> expr = Expression('something.nil', lookup='strict')
cmlenz@442:     >>> expr.evaluate({'something': dict()})
cmlenz@442:     Traceback (most recent call last):
cmlenz@442:         ...
cmlenz@442:     UndefinedError: {} has no member named "nil"
cmlenz@442:     """
cmlenz@442:     def undefined(cls, key, owner=UNDEFINED):
cmlenz@442:         """Raise an ``UndefinedError`` immediately."""
cmlenz@442:         __traceback_hide__ = True
cmlenz@442:         raise UndefinedError(key, owner=owner)
cmlenz@442:     undefined = classmethod(undefined)
cmlenz@442: 
cmlenz@336: 
aronacher@642: class SecurityLookupWrapper(object):
aronacher@642:     """
aronacher@642:     Special class that wraps a lookup so that insecure accesses result
aronacher@642:     in undefined.  Additionally the globals are secured.
aronacher@642:     """
aronacher@642: 
aronacher@642:     def __init__(self, lookup):
aronacher@642:         self._lookup = lookup
aronacher@642: 
aronacher@642:     def __getattr__(self, name):
aronacher@642:         return getattr(self._lookup, name)
aronacher@642: 
aronacher@642:     def globals(self):
aronacher@642:         namespace = self._lookup.globals()
aronacher@642:         namespace.update(
aronacher@642:             _lookup_name=self.lookup_name,
aronacher@642:             _lookup_attr=self.lookup_attr,
aronacher@642:             _lookup_item=self.lookup_item
aronacher@642:         )
aronacher@642:         return namespace
aronacher@642: 
aronacher@642:     def lookup_name(self, data, name):
aronacher@642:         __traceback_hide__ = True
aronacher@642:         if name.startswith('_'):
aronacher@642:             val = self._lookup.undefined(name)
aronacher@642:         else:
aronacher@642:             val = data.get(name, UNDEFINED)
aronacher@642:             if val is UNDEFINED:
aronacher@642:                 val = SECURE_BUILTINS.get(name, val)
aronacher@642:                 if val is UNDEFINED:
aronacher@642:                     val = self._lookup.undefined(name)
aronacher@642:         return val
aronacher@642: 
aronacher@642:     def lookup_attr(self, data, obj, key):
aronacher@642:         __traceback_hide__ = True
aronacher@642:         # XXX: if we weaken this, don't forget to create an
aronacher@642:         # _unsafe_test for sys._active_frames / sys._getframe
aronacher@642:         if key.startswith('_') or self._unsafe_test(obj, key):
aronacher@642:             return self._lookup.undefined(key)
aronacher@642:         return self._lookup.lookup_attr(data, obj, key)
aronacher@642: 
aronacher@642:     def lookup_item(cls, data, obj, key):
aronacher@642:         __traceback_hide__ = True
aronacher@642:         if key.startswith('_') or self._unsafe_test(obj, key):
aronacher@642:             return self._lookup.undefined(key)
aronacher@642:         return self._lookup.lookup_item(data, obj, key)
aronacher@642: 
aronacher@642:     def _unsafe_test(self, obj, key):
aronacher@642:         if isinstance(obj, MethodType):
aronacher@642:             if key in ('im_class', 'im_func', 'im_self'):
aronacher@642:                 return True
aronacher@642:             obj = obj.im_func
aronacher@642:         if isinstance(obj, FunctionType):
aronacher@642:             return key in ('func_closure', 'func_code', 'func_defaults',
aronacher@642:                            'func_dict', 'func_doc', 'func_globals',
aronacher@642:                            'func_name')
aronacher@642:         known_unsafe = getattr(obj, '__genshi_unsafe__', ())
aronacher@642:         return key in known_unsafe
aronacher@642: 
aronacher@642: 
cmlenz@336: def _parse(source, mode='eval'):
cmlenz@601:     source = source.strip()
cmlenz@601:     if mode == 'exec':
cmlenz@601:         lines = [line.expandtabs() for line in source.splitlines()]
cmlenz@601:         first = lines[0]
cmlenz@601:         rest = dedent('\n'.join(lines[1:])).rstrip()
cmlenz@601:         if first.rstrip().endswith(':') and not rest[0].isspace():
cmlenz@601:             rest = '\n'.join(['    %s' % line for line in rest.splitlines()])
cmlenz@601:         source = '\n'.join([first, rest])
cmlenz@336:     if isinstance(source, unicode):
cmlenz@336:         source = '\xef\xbb\xbf' + source.encode('utf-8')
cmlenz@336:     return parse(source, mode)
cmlenz@336: 
cmlenz@604: def _compile(node, source=None, mode='eval', filename=None, lineno=-1,
aronacher@642:              xform=None, secure=False):
cmlenz@604:     if xform is None:
cmlenz@604:         xform = {'eval': ExpressionASTTransformer}.get(mode,
cmlenz@604:                                                        TemplateASTTransformer)
aronacher@642:     tree = xform(secure).visit(node)
cmlenz@336:     if isinstance(filename, unicode):
cmlenz@336:         # unicode file names not allowed for code objects
cmlenz@336:         filename = filename.encode('utf-8', 'replace')
cmlenz@336:     elif not filename:
cmlenz@336:         filename = '<string>'
cmlenz@336:     tree.filename = filename
cmlenz@336:     if lineno <= 0:
cmlenz@336:         lineno = 1
cmlenz@336: 
cmlenz@405:     if mode == 'eval':
cmlenz@405:         gen = ExpressionCodeGenerator(tree)
cmlenz@601:         name = '<Expression %r>' % (source or '?')
cmlenz@405:     else:
cmlenz@405:         gen = ModuleCodeGenerator(tree)
cmlenz@601:         lines = source.splitlines()
cmlenz@601:         extract = lines[0]
cmlenz@601:         if len(lines) > 1:
cmlenz@601:             extract += ' ...'
cmlenz@601:         name = '<Suite %r>' % (extract)
cmlenz@336:     gen.optimized = True
cmlenz@336:     code = gen.getCode()
cmlenz@336: 
cmlenz@336:     # We'd like to just set co_firstlineno, but it's readonly. So we need to
cmlenz@336:     # clone the code object while adjusting the line number
cmlenz@336:     return new.code(0, code.co_nlocals, code.co_stacksize,
cmlenz@336:                     code.co_flags | 0x0040, code.co_code, code.co_consts,
cmlenz@405:                     code.co_names, code.co_varnames, filename, name, lineno,
cmlenz@405:                     code.co_lnotab, (), ())
cmlenz@336: 
cmlenz@336: BUILTINS = __builtin__.__dict__.copy()
cmlenz@442: BUILTINS.update({'Markup': Markup, 'Undefined': Undefined})
aronacher@642: 
aronacher@642: # XXX: if we weaken the rule for global name resultion so that leading
aronacher@642: # underscores are valid we have to add __import__ here.
aronacher@642: UNSAFE_NAMES = ['file', 'open', 'eval', 'locals', 'globals', 'vars',
aronacher@642:                 'help', 'quit', 'exit', 'input', 'raw_input', 'setattr',
aronacher@643:                 'delattr', 'reload', 'compile', 'type']
aronacher@642: 
aronacher@642: # XXX: provide a secure range function
aronacher@642: SECURE_BUILTINS = BUILTINS.copy()
aronacher@642: for _unsafe_name in UNSAFE_NAMES:
aronacher@642:     del SECURE_BUILTINS[_unsafe_name]
aronacher@642: del _unsafe_name
aronacher@643: SECURE_BUILTINS['range'] = safe_range
aronacher@642: 
cmlenz@563: CONSTANTS = frozenset(['False', 'True', 'None', 'NotImplemented', 'Ellipsis'])
cmlenz@336: 
cmlenz@564: 
cmlenz@336: class ASTTransformer(object):
cmlenz@336:     """General purpose base class for AST transformations.
cmlenz@336:     
cmlenz@336:     Every visitor method can be overridden to return an AST node that has been
cmlenz@336:     altered or replaced in some way.
cmlenz@336:     """
cmlenz@336: 
aronacher@642:     def __init__(self, secure):
aronacher@642:         self.secure = secure
aronacher@642: 
cmlenz@357:     def visit(self, node):
cmlenz@405:         if node is None:
cmlenz@405:             return None
cmlenz@473:         if type(node) is tuple:
cmlenz@473:             return tuple([self.visit(n) for n in node])
cmlenz@473:         visitor = getattr(self, 'visit%s' % node.__class__.__name__,
cmlenz@473:                           self._visitDefault)
cmlenz@473:         return visitor(node)
cmlenz@405: 
cmlenz@565:     def _clone(self, node, *args):
cmlenz@565:         lineno = getattr(node, 'lineno', None)
cmlenz@565:         node = node.__class__(*args)
cmlenz@565:         if lineno is not None:
cmlenz@565:             node.lineno = lineno
cmlenz@565:         if isinstance(node, (ast.Class, ast.Function, ast.GenExpr, ast.Lambda)):
cmlenz@565:             node.filename = '<string>' # workaround for bug in pycodegen
cmlenz@565:         return node
cmlenz@565: 
cmlenz@405:     def _visitDefault(self, node):
cmlenz@405:         return node
cmlenz@336: 
cmlenz@357:     def visitExpression(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.node))
cmlenz@336: 
cmlenz@405:     def visitModule(self, node):
cmlenz@565:         return self._clone(node, node.doc, self.visit(node.node))
cmlenz@405: 
cmlenz@405:     def visitStmt(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes])
cmlenz@405: 
cmlenz@405:     # Classes, Functions & Accessors
cmlenz@336: 
cmlenz@357:     def visitCallFunc(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.node),
cmlenz@565:             [self.visit(x) for x in node.args],
cmlenz@565:             node.star_args and self.visit(node.star_args) or None,
cmlenz@565:             node.dstar_args and self.visit(node.dstar_args) or None
cmlenz@565:         )
cmlenz@336: 
cmlenz@405:     def visitClass(self, node):
cmlenz@565:         return self._clone(node, node.name, [self.visit(x) for x in node.bases],
cmlenz@568:             node.doc, self.visit(node.code)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def visitFunction(self, node):
cmlenz@565:         args = []
cmlenz@405:         if hasattr(node, 'decorators'):
cmlenz@565:             args.append(self.visit(node.decorators))
cmlenz@565:         return self._clone(node, *args + [
cmlenz@565:             node.name,
cmlenz@565:             node.argnames,
cmlenz@565:             [self.visit(x) for x in node.defaults],
cmlenz@565:             node.flags,
cmlenz@565:             node.doc,
cmlenz@565:             self.visit(node.code)
cmlenz@565:         ])
cmlenz@336: 
cmlenz@357:     def visitGetattr(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr), node.attrname)
cmlenz@336: 
cmlenz@405:     def visitLambda(self, node):
cmlenz@565:         node = self._clone(node, node.argnames,
cmlenz@565:             [self.visit(x) for x in node.defaults], node.flags,
cmlenz@565:             self.visit(node.code)
cmlenz@565:         )
cmlenz@405:         return node
cmlenz@405: 
cmlenz@357:     def visitSubscript(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr), node.flags,
cmlenz@565:             [self.visit(x) for x in node.subs]
cmlenz@565:         )
cmlenz@336: 
cmlenz@405:     # Statements
cmlenz@405: 
cmlenz@405:     def visitAssert(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.test), self.visit(node.fail))
cmlenz@405: 
cmlenz@405:     def visitAssign(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes],
cmlenz@565:             self.visit(node.expr)
cmlenz@565:         )
cmlenz@405: 
cmlenz@473:     def visitAssAttr(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr), node.attrname,
cmlenz@565:             node.flags
cmlenz@565:         )
cmlenz@473: 
cmlenz@473:     def visitAugAssign(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.node), node.op,
cmlenz@565:             self.visit(node.expr)
cmlenz@565:         )
cmlenz@473: 
cmlenz@405:     def visitDecorators(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes])
cmlenz@405: 
cmlenz@473:     def visitExec(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr), self.visit(node.locals),
cmlenz@565:             self.visit(node.globals)
cmlenz@565:         )
cmlenz@473: 
cmlenz@405:     def visitFor(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.assign), self.visit(node.list),
cmlenz@565:             self.visit(node.body), self.visit(node.else_)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def visitIf(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.tests],
cmlenz@565:             self.visit(node.else_)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def _visitPrint(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes],
cmlenz@565:             self.visit(node.dest)
cmlenz@565:         )
cmlenz@405:     visitPrint = visitPrintnl = _visitPrint
cmlenz@405: 
cmlenz@405:     def visitRaise(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr1), self.visit(node.expr2),
cmlenz@565:             self.visit(node.expr3)
cmlenz@565:         )
cmlenz@405: 
cmlenz@473:     def visitReturn(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.value))
cmlenz@473: 
cmlenz@405:     def visitTryExcept(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.body), self.visit(node.handlers),
cmlenz@565:             self.visit(node.else_)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def visitTryFinally(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.body), self.visit(node.final))
cmlenz@405: 
cmlenz@405:     def visitWhile(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.test), self.visit(node.body),
cmlenz@565:             self.visit(node.else_)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def visitWith(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr),
cmlenz@565:             [self.visit(x) for x in node.vars], self.visit(node.body)
cmlenz@565:         )
cmlenz@405: 
cmlenz@405:     def visitYield(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.value))
cmlenz@405: 
cmlenz@336:     # Operators
cmlenz@336: 
cmlenz@357:     def _visitBoolOp(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes])
mgood@396:     visitAnd = visitOr = visitBitand = visitBitor = visitBitxor = _visitBoolOp
cmlenz@473:     visitAssTuple = visitAssList = _visitBoolOp
cmlenz@336: 
cmlenz@357:     def _visitBinOp(self, node):
cmlenz@565:         return self._clone(node,
cmlenz@565:             (self.visit(node.left), self.visit(node.right))
cmlenz@565:         )
cmlenz@336:     visitAdd = visitSub = _visitBinOp
cmlenz@336:     visitDiv = visitFloorDiv = visitMod = visitMul = visitPower = _visitBinOp
cmlenz@336:     visitLeftShift = visitRightShift = _visitBinOp
cmlenz@336: 
cmlenz@357:     def visitCompare(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr),
cmlenz@565:             [(op, self.visit(n)) for op, n in  node.ops]
cmlenz@565:         )
cmlenz@336: 
cmlenz@357:     def _visitUnaryOp(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr))
cmlenz@336:     visitUnaryAdd = visitUnarySub = visitNot = visitInvert = _visitUnaryOp
cmlenz@405:     visitBackquote = visitDiscard = _visitUnaryOp
cmlenz@336: 
mgood@393:     def visitIfExp(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.test), self.visit(node.then),
cmlenz@565:             self.visit(node.else_)
cmlenz@565:         )
mgood@393: 
cmlenz@336:     # Identifiers, Literals and Comprehensions
cmlenz@336: 
cmlenz@357:     def visitDict(self, node):
cmlenz@565:         return self._clone(node, 
cmlenz@565:             [(self.visit(k), self.visit(v)) for k, v in node.items]
cmlenz@565:         )
cmlenz@336: 
cmlenz@357:     def visitGenExpr(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.code))
cmlenz@336: 
cmlenz@357:     def visitGenExprFor(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.assign), self.visit(node.iter),
cmlenz@565:             [self.visit(x) for x in node.ifs]
cmlenz@565:         )
cmlenz@336: 
cmlenz@357:     def visitGenExprIf(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.test))
cmlenz@336: 
cmlenz@357:     def visitGenExprInner(self, node):
cmlenz@565:         quals = [self.visit(x) for x in node.quals]
cmlenz@565:         return self._clone(node, self.visit(node.expr), quals)
cmlenz@336: 
cmlenz@357:     def visitKeyword(self, node):
cmlenz@565:         return self._clone(node, node.name, self.visit(node.expr))
cmlenz@336: 
cmlenz@357:     def visitList(self, node):
cmlenz@565:         return self._clone(node, [self.visit(n) for n in node.nodes])
cmlenz@336: 
cmlenz@357:     def visitListComp(self, node):
cmlenz@565:         quals = [self.visit(x) for x in node.quals]
cmlenz@565:         return self._clone(node, self.visit(node.expr), quals)
cmlenz@357: 
cmlenz@357:     def visitListCompFor(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.assign), self.visit(node.list),
cmlenz@565:             [self.visit(x) for x in node.ifs]
cmlenz@565:         )
cmlenz@357: 
cmlenz@357:     def visitListCompIf(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.test))
cmlenz@357: 
cmlenz@357:     def visitSlice(self, node):
cmlenz@565:         return self._clone(node, self.visit(node.expr), node.flags,
cmlenz@565:             node.lower and self.visit(node.lower) or None,
cmlenz@565:             node.upper and self.visit(node.upper) or None
cmlenz@565:         )
cmlenz@357: 
cmlenz@357:     def visitSliceobj(self, node):
cmlenz@565:         return self._clone(node, [self.visit(x) for x in node.nodes])
cmlenz@357: 
cmlenz@357:     def visitTuple(self, node):
cmlenz@565:         return self._clone(node, [self.visit(n) for n in node.nodes])
cmlenz@336: 
cmlenz@336: 
cmlenz@405: class TemplateASTTransformer(ASTTransformer):
cmlenz@336:     """Concrete AST transformer that implements the AST transformations needed
cmlenz@405:     for code embedded in templates.
cmlenz@336:     """
cmlenz@336: 
aronacher@642:     def __init__(self, secure):
aronacher@642:         ASTTransformer.__init__(self, secure)
cmlenz@586:         self.locals = [CONSTANTS]
cmlenz@357: 
cmlenz@357:     def visitConst(self, node):
cmlenz@336:         if isinstance(node.value, str):
cmlenz@357:             try: # If the string is ASCII, return a `str` object
cmlenz@357:                 node.value.decode('ascii')
cmlenz@357:             except ValueError: # Otherwise return a `unicode` object
cmlenz@357:                 return ast.Const(node.value.decode('utf-8'))
cmlenz@336:         return node
cmlenz@336: 
cmlenz@357:     def visitAssName(self, node):
cmlenz@586:         if len(self.locals) > 1:
cmlenz@405:             self.locals[-1].add(node.name)
cmlenz@405:         return node
cmlenz@405: 
cmlenz@473:     def visitAugAssign(self, node):
cmlenz@579:         if isinstance(node.node, ast.Name) \
cmlenz@582:                 and node.node.name not in flatten(self.locals):
cmlenz@473:             name = node.node.name
cmlenz@473:             node.node = ast.Subscript(ast.Name('data'), 'OP_APPLY',
cmlenz@473:                                       [ast.Const(name)])
cmlenz@473:             node.expr = self.visit(node.expr)
cmlenz@473:             return ast.If([
cmlenz@473:                 (ast.Compare(ast.Const(name), [('in', ast.Name('data'))]),
cmlenz@473:                  ast.Stmt([node]))],
cmlenz@473:                 ast.Stmt([ast.Raise(ast.CallFunc(ast.Name('UndefinedError'),
cmlenz@473:                                                  [ast.Const(name)]),
cmlenz@473:                                     None, None)]))
cmlenz@473:         else:
cmlenz@473:             return ASTTransformer.visitAugAssign(self, node)
cmlenz@473: 
cmlenz@405:     def visitClass(self, node):
cmlenz@586:         if len(self.locals) > 1:
cmlenz@586:             self.locals[-1].add(node.name)
cmlenz@405:         self.locals.append(set())
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitClass(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@405: 
cmlenz@405:     def visitFor(self, node):
cmlenz@405:         self.locals.append(set())
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitFor(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@405: 
cmlenz@405:     def visitFunction(self, node):
cmlenz@586:         if len(self.locals) > 1:
cmlenz@586:             self.locals[-1].add(node.name)
cmlenz@405:         self.locals.append(set(node.argnames))
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitFunction(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@336: 
cmlenz@357:     def visitGenExpr(self, node):
cmlenz@357:         self.locals.append(set())
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitGenExpr(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@336: 
cmlenz@357:     def visitLambda(self, node):
cmlenz@357:         self.locals.append(set(flatten(node.argnames)))
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitLambda(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@336: 
cmlenz@357:     def visitListComp(self, node):
cmlenz@357:         self.locals.append(set())
cmlenz@565:         try:
cmlenz@565:             return ASTTransformer.visitListComp(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@357: 
cmlenz@357:     def visitName(self, node):
cmlenz@357:         # If the name refers to a local inside a lambda, list comprehension, or
cmlenz@357:         # generator expression, leave it alone
cmlenz@586:         if node.name not in flatten(self.locals):
cmlenz@586:             # Otherwise, translate the name ref into a context lookup
cmlenz@586:             func_args = [ast.Name('data'), ast.Const(node.name)]
cmlenz@586:             node = ast.CallFunc(ast.Name('_lookup_name'), func_args)
cmlenz@586:         return node
cmlenz@336: 
cmlenz@473: 
cmlenz@473: class ExpressionASTTransformer(TemplateASTTransformer):
cmlenz@473:     """Concrete AST transformer that implements the AST transformations needed
cmlenz@473:     for code embedded in templates.
cmlenz@473:     """
cmlenz@473: 
cmlenz@473:     def visitGetattr(self, node):
cmlenz@473:         return ast.CallFunc(ast.Name('_lookup_attr'), [
cmlenz@473:             ast.Name('data'), self.visit(node.expr),
cmlenz@473:             ast.Const(node.attrname)
cmlenz@473:         ])
cmlenz@473: 
cmlenz@357:     def visitSubscript(self, node):
cmlenz@336:         return ast.CallFunc(ast.Name('_lookup_item'), [
cmlenz@357:             ast.Name('data'), self.visit(node.expr),
cmlenz@357:             ast.Tuple([self.visit(sub) for sub in node.subs])
cmlenz@336:         ])