InRough update of the beets lib to 1.0b15

lib/beets/util/functemplate.py

@@ -8,7 +8,7 @@
 # distribute, sublicense, and/or sell copies of the Software, and to
 # permit persons to whom the Software is furnished to do so, subject to
 # the following conditions:
-# 
+#
 # The above copyright notice and this permission notice shall be
 # included in all copies or substantial portions of the Software.
 
@@ -25,7 +25,12 @@ library: unknown symbols are left intact.
 This is sort of like a tiny, horrible degeneration of a real templating
 engine like Jinja2 or Mustache.
 """
+from __future__ import print_function
+
 import re
+import ast
+import dis
+import types
 
 SYMBOL_DELIM = u'$'
 FUNC_DELIM = u'%'
@@ -34,6 +39,9 @@ GROUP_CLOSE = u'}'
 ARG_SEP = u','
 ESCAPE_CHAR = u'$'
 
+VARIABLE_PREFIX = '__var_'
+FUNCTION_PREFIX = '__func_'
+
 class Environment(object):
     """Contains the values and functions to be substituted into a
     template.
@@ -42,6 +50,88 @@ class Environment(object):
         self.values = values
         self.functions = functions
 
+
+# Code generation helpers.
+
+def ex_lvalue(name):
+    """A variable load expression."""
+    return ast.Name(name, ast.Store())
+
+def ex_rvalue(name):
+    """A variable store expression."""
+    return ast.Name(name, ast.Load())
+
+def ex_literal(val):
+    """An int, float, long, bool, string, or None literal with the given
+    value.
+    """
+    if val is None:
+        return ast.Name('None', ast.Load())
+    elif isinstance(val, (int, float, long)):
+        return ast.Num(val)
+    elif isinstance(val, bool):
+        return ast.Name(str(val), ast.Load())
+    elif isinstance(val, basestring):
+        return ast.Str(val)
+    raise TypeError('no literal for {0}'.format(type(val)))
+
+def ex_varassign(name, expr):
+    """Assign an expression into a single variable. The expression may
+    either be an `ast.expr` object or a value to be used as a literal.
+    """
+    if not isinstance(expr, ast.expr):
+        expr = ex_literal(expr)
+    return ast.Assign([ex_lvalue(name)], expr)
+
+def ex_call(func, args):
+    """A function-call expression with only positional parameters. The
+    function may be an expression or the name of a function. Each
+    argument may be an expression or a value to be used as a literal.
+    """
+    if isinstance(func, basestring):
+        func = ex_rvalue(func)
+
+    args = list(args)
+    for i in range(len(args)):
+        if not isinstance(args[i], ast.expr):
+            args[i] = ex_literal(args[i])
+
+    return ast.Call(func, args, [], None, None)
+
+def compile_func(arg_names, statements, name='_the_func', debug=False):
+    """Compile a list of statements as the body of a function and return
+    the resulting Python function. If `debug`, then print out the
+    bytecode of the compiled function.
+    """
+    func_def = ast.FunctionDef(
+        name,
+        ast.arguments(
+            [ast.Name(n, ast.Param()) for n in arg_names],
+            None, None,
+            [ex_literal(None) for _ in arg_names],
+        ),
+        statements,
+        [],
+    )
+    mod = ast.Module([func_def])
+    ast.fix_missing_locations(mod)
+
+    prog = compile(mod, '<generated>', 'exec')
+
+    # Debug: show bytecode.
+    if debug:
+        dis.dis(prog)
+        for const in prog.co_consts:
+            if isinstance(const, types.CodeType):
+                dis.dis(const)
+
+    the_locals = {}
+    exec prog in {}, the_locals
+    return the_locals[name]
+
+
+# AST nodes for the template language.
+
 class Symbol(object):
     """A variable-substitution symbol in a template."""
     def __init__(self, ident, original):
@@ -62,6 +152,11 @@ class Symbol(object):
             # Keep original text.
             return self.original
 
+    def translate(self):
+        """Compile the variable lookup."""
+        expr = ex_rvalue(VARIABLE_PREFIX + self.ident.encode('utf8'))
+        return [expr], set([self.ident.encode('utf8')]), set()
+
 class Call(object):
     """A function call in a template."""
     def __init__(self, ident, args, original):
@@ -81,7 +176,7 @@ class Call(object):
             arg_vals = [expr.evaluate(env) for expr in self.args]
             try:
                 out = env.functions[self.ident](*arg_vals)
-            except Exception, exc:
+            except Exception as exc:
                 # Function raised exception! Maybe inlining the name of
                 # the exception will help debug.
                 return u'<%s>' % unicode(exc)
@@ -89,6 +184,36 @@ class Call(object):
         else:
             return self.original
 
+    def translate(self):
+        """Compile the function call."""
+        varnames = set()
+        funcnames = set([self.ident.encode('utf8')])
+
+        arg_exprs = []
+        for arg in self.args:
+            subexprs, subvars, subfuncs = arg.translate()
+            varnames.update(subvars)
+            funcnames.update(subfuncs)
+
+            # Create a subexpression that joins the result components of
+            # the arguments.
+            arg_exprs.append(ex_call(
+                ast.Attribute(ex_literal(u''), 'join', ast.Load()),
+                [ex_call(
+                    'map',
+                    [
+                        ex_rvalue('unicode'),
+                        ast.List(subexprs, ast.Load()),
+                    ]
+                )],
+            ))
+
+        subexpr_call = ex_call(
+            FUNCTION_PREFIX + self.ident.encode('utf8'),
+            arg_exprs
+        )
+        return [subexpr_call], varnames, funcnames
+
 class Expression(object):
     """Top-level template construct: contains a list of text blobs,
     Symbols, and Calls.
@@ -111,6 +236,26 @@ class Expression(object):
                 out.append(part.evaluate(env))
         return u''.join(map(unicode, out))
 
+    def translate(self):
+        """Compile the expression to a list of Python AST expressions, a
+        set of variable names used, and a set of function names.
+        """
+        expressions = []
+        varnames = set()
+        funcnames = set()
+        for part in self.parts:
+            if isinstance(part, basestring):
+                expressions.append(ex_literal(part))
+            else:
+                e, v, f = part.translate()
+                expressions.extend(e)
+                varnames.update(v)
+                funcnames.update(f)
+        return expressions, varnames, funcnames
+
+
+# Parser.
+
 class ParseError(Exception):
     pass
 
@@ -266,7 +411,7 @@ class Parser(object):
             # No function name.
             self.parts.append(FUNC_DELIM)
             return
-        
+
         if self.pos >= len(self.string):
             # Identifier terminates string.
             self.parts.append(self.string[start_pos:self.pos])
@@ -304,7 +449,7 @@ class Parser(object):
 
             # Extract and advance past the parsed expression.
             expressions.append(Expression(subparser.parts))
-            self.pos += subparser.pos 
+            self.pos += subparser.pos
 
             if self.pos >= len(self.string) or \
                self.string[self.pos] == GROUP_CLOSE:
@@ -340,14 +485,74 @@ def _parse(template):
         parts.append(remainder)
     return Expression(parts)
 
+
+# External interface.
+
 class Template(object):
     """A string template, including text, Symbols, and Calls.
     """
     def __init__(self, template):
         self.expr = _parse(template)
         self.original = template
+        self.compiled = self.translate()
+
+    def interpret(self, values={}, functions={}):
+        """Like `substitute`, but forces the interpreter (rather than
+        the compiled version) to be used. The interpreter includes
+        exception-handling code for missing variables and buggy template
+        functions but is much slower.
+        """
+        return self.expr.evaluate(Environment(values, functions))
 
     def substitute(self, values={}, functions={}):
         """Evaluate the template given the values and functions.
         """
-        return self.expr.evaluate(Environment(values, functions))
+        try:
+            res = self.compiled(values, functions)
+        except:  # Handle any exceptions thrown by compiled version.
+            res = self.interpret(values, functions)
+        return res
+
+    def translate(self):
+        """Compile the template to a Python function."""
+        expressions, varnames, funcnames = self.expr.translate()
+
+        argnames = []
+        for varname in varnames:
+            argnames.append(VARIABLE_PREFIX.encode('utf8') + varname)
+        for funcname in funcnames:
+            argnames.append(FUNCTION_PREFIX.encode('utf8') + funcname)
+
+        func = compile_func(
+            argnames,
+            [ast.Return(ast.List(expressions, ast.Load()))],
+        )
+
+        def wrapper_func(values={}, functions={}):
+            args = {}
+            for varname in varnames:
+                args[VARIABLE_PREFIX + varname] = values[varname]
+            for funcname in funcnames:
+                args[FUNCTION_PREFIX + funcname] = functions[funcname]
+            parts = func(**args)
+            return u''.join(parts)
+
+        return wrapper_func
+
+
+# Performance tests.
+
+if __name__ == '__main__':
+    import timeit
+    _tmpl = Template(u'foo $bar %baz{foozle $bar barzle} $bar')
+    _vars = {'bar': 'qux'}
+    _funcs = {'baz': unicode.upper}
+    interp_time = timeit.timeit('_tmpl.interpret(_vars, _funcs)',
+                                'from __main__ import _tmpl, _vars, _funcs',
+                                number=10000)
+    print(interp_time)
+    comp_time = timeit.timeit('_tmpl.substitute(_vars, _funcs)',
+                              'from __main__ import _tmpl, _vars, _funcs',
+                              number=10000)
+    print(comp_time)
+    print('Speedup:', interp_time / comp_time)