cmlenz@336: # -*- coding: utf-8 -*-
cmlenz@336: #
cmlenz@854: # Copyright (C) 2006-2009 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@794: 
cmlenz@601: from textwrap import dedent
cmlenz@793: from types import CodeType
cmlenz@357: 
cmlenz@401: from genshi.core import Markup
cmlenz@794: from genshi.template.astutil import ASTTransformer, ASTCodeGenerator, \
cmlenz@794:                                     _ast, parse
cmlenz@418: from genshi.template.base import TemplateRuntimeError
cmlenz@357: from genshi.util import flatten
cmlenz@336: 
cmlenz@442: __all__ = ['Code', 'Expression', 'Suite', 'LenientLookup', 'StrictLookup',
cmlenz@442:            'Undefined', 'UndefinedError']
cmlenz@425: __docformat__ = 'restructuredtext en'
cmlenz@336: 
cmlenz@794: 
cmlenz@731: # Check for a Python 2.4 bug in the eval loop
cmlenz@739: has_star_import_bug = False
cmlenz@731: try:
cmlenz@736:     class _FakeMapping(object):
cmlenz@731:         __getitem__ = __setitem__ = lambda *a: None
cmlenz@731:     exec 'from sys import *' in {}, _FakeMapping()
cmlenz@739: except SystemError:
cmlenz@731:     has_star_import_bug = True
cmlenz@731: del _FakeMapping
cmlenz@731: 
cmlenz@794: 
cmlenz@731: def _star_import_patch(mapping, modname):
cmlenz@731:     """This function is used as helper if a Python version with a broken
cmlenz@731:     star-import opcode is in use.
cmlenz@731:     """
cmlenz@731:     module = __import__(modname, None, None, ['__all__'])
cmlenz@731:     if hasattr(module, '__all__'):
cmlenz@731:         members = module.__all__
cmlenz@731:     else:
cmlenz@731:         members = [x for x in module.__dict__ if not x.startswith('_')]
athomas@733:     mapping.update([(name, getattr(module, name)) for name in members])
cmlenz@731: 
cmlenz@336: 
cmlenz@405: class Code(object):
cmlenz@405:     """Abstract base class for the `Expression` and `Suite` classes."""
cmlenz@565:     __slots__ = ['source', 'code', 'ast', '_globals']
cmlenz@405: 
cmlenz@606:     def __init__(self, source, filename=None, lineno=-1, lookup='strict',
cmlenz@604:                  xform=None):
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
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@794:             assert isinstance(source, _ast.AST), \
cmlenz@601:                 'Expected string or AST node, but got %r' % source
cmlenz@405:             self.source = '?'
cmlenz@405:             if self.mode == 'eval':
cmlenz@794:                 node = _ast.Expression()
cmlenz@794:                 node.body = source
cmlenz@405:             else:
cmlenz@794:                 node = _ast.Module()
cmlenz@794:                 node.body = [source]
cmlenz@405: 
cmlenz@565:         self.ast = node
cmlenz@405:         self.code = _compile(node, self.source, mode=self.mode,
cmlenz@604:                              filename=filename, lineno=lineno, xform=xform)
cmlenz@442:         if lookup is None:
cmlenz@442:             lookup = LenientLookup
cmlenz@442:         elif isinstance(lookup, basestring):
cmlenz@442:             lookup = {'lenient': LenientLookup, 'strict': StrictLookup}[lookup]
cmlenz@654:         self._globals = lookup.globals
cmlenz@405: 
cmlenz@715:     def __getstate__(self):
cmlenz@715:         state = {'source': self.source, 'ast': self.ast,
cmlenz@715:                  'lookup': self._globals.im_self}
cmlenz@715:         c = self.code
cmlenz@715:         state['code'] = (c.co_nlocals, c.co_stacksize, c.co_flags, c.co_code,
cmlenz@715:                          c.co_consts, c.co_names, c.co_varnames, c.co_filename,
cmlenz@715:                          c.co_name, c.co_firstlineno, c.co_lnotab, (), ())
cmlenz@715:         return state
cmlenz@715: 
cmlenz@715:     def __setstate__(self, state):
cmlenz@715:         self.source = state['source']
cmlenz@715:         self.ast = state['ast']
cmlenz@794:         self.code = CodeType(0, *state['code'])
cmlenz@715:         self._globals = state['lookup'].globals
cmlenz@715: 
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@860:         return '%s(%r)' % (type(self).__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@691:         _globals = self._globals(data)
cmlenz@682:         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@691:         _globals = self._globals(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@853:     >>> 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@860:         return '<%s %r>' % (type(self).__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@878:     # Hack around some behavior introduced in Python 2.6.2
cmlenz@878:     # http://genshi.edgewall.org/ticket/324
cmlenz@878:     __length_hint__ = None
cmlenz@878: 
cmlenz@442: 
cmlenz@442: class LookupBase(object):
cmlenz@442:     """Abstract base class for variable lookup implementations."""
cmlenz@442: 
cmlenz@822:     @classmethod
cmlenz@691:     def globals(cls, data):
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@691:             '__data__': data,
cmlenz@442:             '_lookup_name': cls.lookup_name,
cmlenz@442:             '_lookup_attr': cls.lookup_attr,
cmlenz@579:             '_lookup_item': cls.lookup_item,
cmlenz@731:             '_star_import_patch': _star_import_patch,
cmlenz@691:             'UndefinedError': UndefinedError,
cmlenz@442:         }
cmlenz@442: 
cmlenz@822:     @classmethod
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: 
cmlenz@822:     @classmethod
cmlenz@691:     def lookup_attr(cls, obj, key):
cmlenz@442:         __traceback_hide__ = True
cmlenz@702:         try:
cmlenz@702:             val = getattr(obj, key)
cmlenz@702:         except AttributeError:
cmlenz@702:             if hasattr(obj.__class__, key):
cmlenz@702:                 raise
cmlenz@702:             else:
cmlenz@702:                 try:
cmlenz@702:                     val = obj[key]
cmlenz@702:                 except (KeyError, TypeError):
cmlenz@702:                     val = cls.undefined(key, owner=obj)
cmlenz@569:         return val
cmlenz@442: 
cmlenz@822:     @classmethod
cmlenz@691:     def lookup_item(cls, 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: 
cmlenz@822:     @classmethod
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: 
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@822: 
cmlenz@822:     @classmethod
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: 
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@822: 
cmlenz@822:     @classmethod
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: 
cmlenz@336: 
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()]
aflett@716:         if lines:
aflett@716:             first = lines[0]
aflett@716:             rest = dedent('\n'.join(lines[1:])).rstrip()
aflett@716:             if first.rstrip().endswith(':') and not rest[0].isspace():
aflett@716:                 rest = '\n'.join(['    %s' % line for line in rest.splitlines()])
aflett@716:             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@794: 
cmlenz@604: def _compile(node, source=None, mode='eval', filename=None, lineno=-1,
cmlenz@604:              xform=None):
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:     if lineno <= 0:
cmlenz@336:         lineno = 1
cmlenz@336: 
cmlenz@794:     if xform is None:
cmlenz@794:         xform = {
cmlenz@794:             'eval': ExpressionASTTransformer
cmlenz@794:         }.get(mode, TemplateASTTransformer)
cmlenz@794:     tree = xform().visit(node)
cmlenz@794: 
cmlenz@405:     if mode == 'eval':
cmlenz@601:         name = '<Expression %r>' % (source or '?')
cmlenz@405:     else:
cmlenz@601:         lines = source.splitlines()
aflett@716:         if not lines:
aflett@716:             extract = ''
aflett@716:         else:
aflett@716:             extract = lines[0]
cmlenz@601:         if len(lines) > 1:
cmlenz@601:             extract += ' ...'
cmlenz@601:         name = '<Suite %r>' % (extract)
cmlenz@794:     new_source = ASTCodeGenerator(tree).code
cmlenz@794:     code = compile(new_source, filename, mode)
cmlenz@336: 
cmlenz@794:     try:
cmlenz@794:         # We'd like to just set co_firstlineno, but it's readonly. So we need
cmlenz@794:         # to clone the code object while adjusting the line number
cmlenz@794:         return CodeType(0, code.co_nlocals, code.co_stacksize,
cmlenz@794:                         code.co_flags | 0x0040, code.co_code, code.co_consts,
cmlenz@794:                         code.co_names, code.co_varnames, filename, name,
cmlenz@794:                         lineno, code.co_lnotab, (), ())
cmlenz@794:     except RuntimeError:
cmlenz@794:         return code
cmlenz@794: 
cmlenz@794: 
cmlenz@794: def _new(class_, *args, **kwargs):
cmlenz@794:     ret = class_()
cmlenz@794:     for attr, value in zip(ret._fields, args):
cmlenz@794:         if attr in kwargs:
cmlenz@794:             raise ValueError('Field set both in args and kwargs')
cmlenz@794:         setattr(ret, attr, value)
cmlenz@794:     for attr, value in kwargs:
cmlenz@794:         setattr(ret, attr, value)
cmlenz@794:     return ret
cmlenz@794: 
cmlenz@336: 
cmlenz@336: BUILTINS = __builtin__.__dict__.copy()
cmlenz@442: BUILTINS.update({'Markup': Markup, 'Undefined': Undefined})
cmlenz@563: CONSTANTS = frozenset(['False', 'True', 'None', 'NotImplemented', 'Ellipsis'])
cmlenz@336: 
cmlenz@564: 
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: 
cmlenz@357:     def __init__(self):
cmlenz@586:         self.locals = [CONSTANTS]
cmlenz@357: 
cmlenz@794:     def _extract_names(self, node):
cmlenz@806:         names = set()
cmlenz@794:         def _process(node):
cmlenz@794:             if isinstance(node, _ast.Name):
cmlenz@806:                 names.add(node.id)
cmlenz@806:             elif isinstance(node, _ast.alias):
cmlenz@806:                 names.add(node.asname or node.name)
cmlenz@794:             elif isinstance(node, _ast.Tuple):
cmlenz@794:                 for elt in node.elts:
cmlenz@896:                     _process(elt)
cmlenz@806:         if hasattr(node, 'args'):
cmlenz@806:             for arg in node.args:
cmlenz@806:                 _process(arg)
cmlenz@806:             if hasattr(node, 'vararg'):
cmlenz@806:                 names.add(node.vararg)
cmlenz@806:             if hasattr(node, 'kwarg'):
cmlenz@806:                 names.add(node.kwarg)
cmlenz@806:         elif hasattr(node, 'names'):
cmlenz@806:             for elt in node.names:
cmlenz@806:                 _process(elt)
cmlenz@806:         return names
cmlenz@794: 
cmlenz@794:     def visit_Str(self, node):
cmlenz@794:         if isinstance(node.s, str):
cmlenz@357:             try: # If the string is ASCII, return a `str` object
cmlenz@794:                 node.s.decode('ascii')
cmlenz@357:             except ValueError: # Otherwise return a `unicode` object
cmlenz@794:                 return _new(_ast.Str, node.s.decode('utf-8'))
cmlenz@405:         return node
cmlenz@405: 
cmlenz@794:     def visit_ClassDef(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@794:             return ASTTransformer.visit_ClassDef(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@336: 
cmlenz@806:     def visit_Import(self, node):
cmlenz@806:         if len(self.locals) > 1:
cmlenz@806:             self.locals[-1].update(self._extract_names(node))
cmlenz@806:         return ASTTransformer.visit_Import(self, node)
cmlenz@806: 
cmlenz@794:     def visit_ImportFrom(self, node):
cmlenz@806:         if [a.name for a in node.names] == ['*']:
cmlenz@806:             if has_star_import_bug:
cmlenz@806:                 # This is a Python 2.4 bug. Only if we have a broken Python
cmlenz@806:                 # version do we need to apply this hack
cmlenz@806:                 node = _new(_ast.Expr, _new(_ast.Call,
cmlenz@806:                     _new(_ast.Name, '_star_import_patch'), [
cmlenz@806:                         _new(_ast.Name, '__data__'),
cmlenz@806:                         _new(_ast.Str, node.module)
cmlenz@806:                     ], (), ()))
cmlenz@794:             return node
cmlenz@806:         if len(self.locals) > 1:
cmlenz@806:             self.locals[-1].update(self._extract_names(node))
cmlenz@806:         return ASTTransformer.visit_ImportFrom(self, node)
cmlenz@794: 
cmlenz@794:     def visit_FunctionDef(self, node):
cmlenz@794:         if len(self.locals) > 1:
cmlenz@794:             self.locals[-1].add(node.name)
cmlenz@794: 
cmlenz@794:         self.locals.append(self._extract_names(node.args))
cmlenz@565:         try:
cmlenz@794:             return ASTTransformer.visit_FunctionDef(self, node)
cmlenz@565:         finally:
cmlenz@565:             self.locals.pop()
cmlenz@357: 
cmlenz@794:     # GeneratorExp(expr elt, comprehension* generators)
cmlenz@794:     def visit_GeneratorExp(self, node):
cmlenz@794:         gens = []
mgood@845:         for generator in node.generators:
cmlenz@794:             # comprehension = (expr target, expr iter, expr* ifs)
cmlenz@794:             self.locals.append(set())
cmlenz@794:             gen = _new(_ast.comprehension, self.visit(generator.target),
mgood@845:                        self.visit(generator.iter),
mgood@845:                        [self.visit(if_) for if_ in generator.ifs])
cmlenz@794:             gens.append(gen)
cmlenz@794: 
cmlenz@794:         # use node.__class__ to make it reusable as ListComp
cmlenz@794:         ret = _new(node.__class__, self.visit(node.elt), gens)
cmlenz@794:         #delete inserted locals
cmlenz@794:         del self.locals[-len(node.generators):]
cmlenz@794:         return ret
cmlenz@794: 
cmlenz@794:     # ListComp(expr elt, comprehension* generators)
cmlenz@794:     visit_ListComp = visit_GeneratorExp
cmlenz@794: 
cmlenz@794:     def visit_Lambda(self, node):
cmlenz@794:         self.locals.append(self._extract_names(node.args))
cmlenz@794:         try:
cmlenz@794:             return ASTTransformer.visit_Lambda(self, node)
cmlenz@794:         finally:
cmlenz@794:             self.locals.pop()
cmlenz@794: 
cmlenz@794:     def visit_Name(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@794:         if isinstance(node.ctx, _ast.Load) and \
cmlenz@794:                 node.id not in flatten(self.locals):
cmlenz@586:             # Otherwise, translate the name ref into a context lookup
cmlenz@794:             name = _new(_ast.Name, '_lookup_name', _ast.Load())
cmlenz@794:             namearg = _new(_ast.Name, '__data__', _ast.Load())
cmlenz@794:             strarg = _new(_ast.Str, node.id)
cmlenz@794:             node = _new(_ast.Call, name, [namearg, strarg], [])
cmlenz@794:         elif isinstance(node.ctx, _ast.Store):
cmlenz@794:             if len(self.locals) > 1:
cmlenz@794:                 self.locals[-1].add(node.id)
cmlenz@794: 
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@794:     def visit_Attribute(self, node):
cmlenz@794:         if not isinstance(node.ctx, _ast.Load):
cmlenz@794:             return ASTTransformer.visit_Attribute(self, node)
cmlenz@473: 
cmlenz@794:         func = _new(_ast.Name, '_lookup_attr', _ast.Load())
cmlenz@794:         args = [self.visit(node.value), _new(_ast.Str, node.attr)]
cmlenz@794:         return _new(_ast.Call, func, args, [])
cmlenz@794: 
cmlenz@794:     def visit_Subscript(self, node):
cmlenz@794:         if not isinstance(node.ctx, _ast.Load) or \
cmlenz@794:                 not isinstance(node.slice, _ast.Index):
cmlenz@794:             return ASTTransformer.visit_Subscript(self, node)
cmlenz@794: 
cmlenz@794:         func = _new(_ast.Name, '_lookup_item', _ast.Load())
cmlenz@794:         args = [
cmlenz@794:             self.visit(node.value),
cmlenz@794:             _new(_ast.Tuple, (self.visit(node.slice.value),), _ast.Load())
cmlenz@794:         ]
cmlenz@794:         return _new(_ast.Call, func, args, [])