Merge pull request #510 from efficiosoft/vyos_template_rewrite

vyos.template: T2720: Rework vyos.template Python library

2 files changed, 119 insertions, 67 deletions

diff --git a/python/vyos/template.py b/python/vyos/template.py
index d9b0c749d..c88ab04a0 100644
--- a/python/vyos/template.py
+++ b/python/vyos/template.py
@@ -1,4 +1,4 @@
-# Copyright 2019 VyOS maintainers and contributors <maintainers@vyos.io>
+# Copyright 2019-2020 VyOS maintainers and contributors <maintainers@vyos.io>
 #
 # This library is free software; you can redistribute it and/or
 # modify it under the terms of the GNU Lesser General Public
@@ -13,78 +13,131 @@
 # You should have received a copy of the GNU Lesser General Public
 # License along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
+import functools
 import os
 
+from ipaddress import ip_network
 from jinja2 import Environment
 from jinja2 import FileSystemLoader
 from vyos.defaults import directories
 from vyos.util import chmod, chown, makedir
 
-# reuse the same Environment to improve performance
-_templates_env = {
-    False: Environment(loader=FileSystemLoader(directories['templates'])),
-    True:  Environment(loader=FileSystemLoader(directories['templates']), trim_blocks=True),
-}
-_templates_mem = {
-    False: {},
-    True: {},
-}
 
-def vyos_address_from_cidr(text):
-    """ Take an IPv4/IPv6 CIDR prefix and convert the network to an "address".
-    Example:
-    192.0.2.0/24 -> 192.0.2.0, 2001:db8::/48 -> 2001:db8::
-    """
-    from ipaddress import ip_network
-    return ip_network(text).network_address
+# Holds template filters registered via register_filter()
+_FILTERS = {}
 
-def vyos_netmask_from_cidr(text):
-    """ Take an IPv4/IPv6 CIDR prefix and convert the prefix length to a "subnet mask".
-    Example:
-    192.0.2.0/24 -> 255.255.255.0, 2001:db8::/48 -> ffff:ffff:ffff::
-    """
-    from ipaddress import ip_network
-    return ip_network(text).netmask
 
-def render(destination, template, content, trim_blocks=False, formater=None, permission=None, user=None, group=None):
+# reuse Environments with identical trim_blocks setting to improve performance
+@functools.lru_cache(maxsize=2)
+def _get_environment(trim_blocks):
+    env = Environment(
+        # Don't check if template files were modified upon re-rendering
+        auto_reload=False,
+        # Cache up to this number of templates for quick re-rendering
+        cache_size=100,
+        loader=FileSystemLoader(directories["templates"]),
+        trim_blocks=trim_blocks,
+    )
+    env.filters.update(_FILTERS)
+    return env
+
+
+def register_filter(name, func=None):
+    """Register a function to be available as filter in templates under given name.
+
+    It can also be used as a decorator, see below in this module for examples.
+
+    :raise RuntimeError:
+        when trying to register a filter after a template has been rendered already
+    :raise ValueError: when trying to register a name which was taken already
     """
-    render a template from the template directory, it will raise on any errors
-    destination: the file where the rendered template must be saved
-    template: the path to the template relative to the template folder
-    content: the dictionary to use to render the template
-
-    This classes cache the renderer, so rendering the same file multiple time
-    does not cause as too much overhead. If use everywhere, it could be changed
-    and load the template from python environement variables from an import
-    python module generated when the debian package is build
-    (recovering the load time and overhead caused by having the file out of the code)
+    if func is None:
+        return functools.partial(register_filter, name)
+    if _get_environment.cache_info().currsize:
+        raise RuntimeError(
+            "Filters can only be registered before rendering the first template"
+        )
+    if name in _FILTERS:
+        raise ValueError(f"A filter with name {name!r} was registered already")
+    _FILTERS[name] = func
+    return func
+
+
+def render_to_string(template, content, trim_blocks=False, formater=None):
+    """Render a template from the template directory, raise on any errors.
+
+    :param template: the path to the template relative to the template folder
+    :param content: the dictionary of variables to put into rendering context
+    :param trim_blocks: controls the trim_blocks jinja2 feature
+    :param formater:
+        if given, it has to be a callable the rendered string is passed through
+
+    The parsed template files are cached, so rendering the same file multiple times
+    does not cause as too much overhead.
+    If used everywhere, it could be changed to load the template from Python
+    environment variables from an importable Python module generated when the Debian
+    package is build (recovering the load time and overhead caused by having the
+    file out of the code).
     """
+    template = _get_environment(bool(trim_blocks)).get_template(template)
+    rendered = template.render(content)
+    if formater is not None:
+        rendered = formater(rendered)
+    return rendered
+
+
+def render(
+    destination,
+    template,
+    content,
+    trim_blocks=False,
+    formater=None,
+    permission=None,
+    user=None,
+    group=None,
+):
+    """Render a template from the template directory to a file, raise on any errors.
 
-    # Create the directory if it does not exists
+    :param destination: path to the file to save the rendered template in
+    :param permission: permission bitmask to set for the output file
+    :param user: user to own the output file
+    :param group: group to own the output file
+
+    All other parameters are as for :func:`render_to_string`.
+    """
+    # Create the directory if it does not exist
     folder = os.path.dirname(destination)
     makedir(folder, user, group)
 
-    # Setup a renderer for the given template
-    # This is cached and re-used for performance
-    if template not in _templates_mem[trim_blocks]:
-        _env = _templates_env[trim_blocks]
-        _env.filters['address_from_cidr'] = vyos_address_from_cidr
-        _env.filters['netmask_from_cidr'] = vyos_netmask_from_cidr
-        _templates_mem[trim_blocks][template] = _env.get_template(template)
+    # As we are opening the file with 'w', we are performing the rendering before
+    # calling open() to not accidentally erase the file if rendering fails
+    rendered = render_to_string(template, content, trim_blocks, formater)
 
-    template = _templates_mem[trim_blocks][template]
+    # Write to file
+    with open(destination, "w") as file:
+        chmod(file.fileno(), permission)
+        chown(file.fileno(), user, group)
+        file.write(rendered)
 
-    # As we are opening the file with 'w', we are performing the rendering
-    # before calling open() to not accidentally erase the file if the
-    # templating fails
-    content = template.render(content)
 
-    if formater:
-        content = formater(content)
+##################################
+# Custom template filters follow #
+##################################
 
-    # Write client config file
-    with open(destination, 'w') as f:
-        f.write(content)
 
-    chmod(destination, permission)
-    chown(destination, user, group)
+@register_filter("address_from_cidr")
+def vyos_address_from_cidr(text):
+    """ Take an IPv4/IPv6 CIDR prefix and convert the network to an "address".
+    Example:
+    192.0.2.0/24 -> 192.0.2.0, 2001:db8::/48 -> 2001:db8::
+    """
+    return ip_network(text).network_address
+
+
+@register_filter("netmask_from_cidr")
+def vyos_netmask_from_cidr(text):
+    """ Take an IPv4/IPv6 CIDR prefix and convert the prefix length to a "subnet mask".
+    Example:
+    192.0.2.0/24 -> 255.255.255.0, 2001:db8::/48 -> ffff:ffff:ffff::
+    """
+    return ip_network(text).netmask
diff --git a/python/vyos/util.py b/python/vyos/util.py
index c07fef599..d27a8a3e7 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -240,7 +240,8 @@ def chown(path, user, group):
     if user is None or group is None:
         return False
 
-    if not os.path.exists(path):
+    # path may also be an open file descriptor
+    if not isinstance(path, int) and not os.path.exists(path):
         return False
 
     uid = getpwnam(user).pw_uid
@@ -250,7 +251,8 @@ def chown(path, user, group):
 
 
 def chmod(path, bitmask):
-    if not os.path.exists(path):
+    # path may also be an open file descriptor
+    if not isinstance(path, int) and not os.path.exists(path):
         return
     if bitmask is None:
         return
@@ -261,28 +263,25 @@ def chmod_600(path):
     """ make file only read/writable by owner """
     from stat import S_IRUSR, S_IWUSR
 
-    if os.path.exists(path):
-        bitmask = S_IRUSR | S_IWUSR
-        os.chmod(path, bitmask)
+    bitmask = S_IRUSR | S_IWUSR
+    chmod(path, bitmask)
 
 
 def chmod_750(path):
     """ make file/directory only executable to user and group """
     from stat import S_IRUSR, S_IWUSR, S_IXUSR, S_IRGRP, S_IXGRP
 
-    if os.path.exists(path):
-        bitmask = S_IRUSR | S_IWUSR | S_IXUSR | S_IRGRP | S_IXGRP
-        os.chmod(path, bitmask)
+    bitmask = S_IRUSR | S_IWUSR | S_IXUSR | S_IRGRP | S_IXGRP
+    chmod(path, bitmask)
 
 
 def chmod_755(path):
     """ make file executable by all """
     from stat import S_IRUSR, S_IWUSR, S_IXUSR, S_IRGRP, S_IXGRP, S_IROTH, S_IXOTH
 
-    if os.path.exists(path):
-        bitmask = S_IRUSR | S_IWUSR | S_IXUSR | S_IRGRP | S_IXGRP | \
-                  S_IROTH | S_IXOTH
-        os.chmod(path, bitmask)
+    bitmask = S_IRUSR | S_IWUSR | S_IXUSR | S_IRGRP | S_IXGRP | \
+              S_IROTH | S_IXOTH
+    chmod(path, bitmask)
 
 
 def makedir(path, user=None, group=None):