output.MakeFilename accepts prefix and suffix. They can be used to cr…

…eate future parts of file names before they are created (e.g. suffix="_log"). Code improved.

lena/flow/read_root_file.py

@@ -1,5 +1,6 @@
-import inspect
+# needs ROOT installed
 import copy
+import inspect
 import sys
 
 import ROOT

lena/output/make_filename.py

@@ -9,99 +9,164 @@ class MakeFilename(object):
     """Make file name, file extension and directory name."""
 
     def __init__(self, filename=None, dirname=None, fileext=None,
-                 overwrite=False):
-        """A single argument can be a string, which will be
+                 prefix=None, suffix=None, overwrite=False):
+        """*filename* is a string, which will be
         used as a file name without extension
         (but it can contain a relative path).
-        The string can contain arguments enclosed in double braces.
+        The string can contain formatting arguments
+        enclosed in double braces.
         These arguments will be filled from context
         during :meth:`__call__`. Example:
 
             MakeFilename("{{variable.type}}/{{variable.name}}")
 
-        By default, values with *context.output* already containing
-        *filename* are not updated (returned unchanged).
+        *dirname* and *fileext* set directory name and file extension.
+        They are treated similarly to *filename* in most aspects.
+
+        It is possible to "postpone" file name creation, but to provide
+        a part of a future file name through *prefix* or *suffix*.
+        They will be appended to file name during its creation.
+        Existing file names are not affected.
+        It is not allowed to use *prefix* or *suffix* if *filename*
+        argument is given.
+
+        For example, if one creates logarithmic plots, but complete
+        file names will be made later,
+        one may use *MakeFilename(suffix="_log")*.
+
+        All these arguments must be strings,
+        otherwise :exc:`.LenaTypeError` is raised.
+        They may all contain formatting arguments.
+
+        By default, values with *context.output*
+        already containing *filename*, *dirname* or *fileext*
+        are not updated (pass unaltered).
         This can be changed using a keyword argument *overwrite*.
-        If context doesn't contain all necessary keys for formatting,
-        it will not be updated.
         For more options, use :class:`.lena.context.UpdateContext`.
 
-        Other allowed keywords are *filename*, *dirname*,
-        *fileext*. Their value must be a string,
-        otherwise :exc:`.LenaTypeError` is raised.
-        At least one of the must be present,
+        At least one argument must be present,
         or :exc:`.LenaTypeError` will be raised.
-        If a simple check finds unbalanced
-        or single braces instead of double,
-        :exc:`.LenaValueError` is raised.
         """
         self._overwrite = bool(overwrite)
 
-        args_supplied = (
-            filename is not None or dirname is not None
-            or fileext is not None)
-        if not args_supplied:
-            # for wrong initialization there must be a TypeError,
-            # not a ValueError. As if it was an obligatory argument.
-            raise lena.core.LenaTypeError(
-                "MakeFilename must be initialized with at least "
-                "one of filename, dirname, fileext"
-            )
+        args = [filename, dirname, fileext, prefix, suffix]
+        args_names = "filename, dirname, fileext, prefix, suffix"
 
-        self._filename = None
-        self._dirname = None
-        self._fileext = None
+        for arg in args:
+            # we don't check Python 2 basestrings from here.
+            if arg is not None and not isinstance(arg, str):
+                raise lena.core.LenaTypeError(
+                    "arguments must be strings, {} "
+                    "provided".format(arg)
+                )
 
-        # todo: rename to filename etc?
-        arg_error = "{} must be a string, {} provided"
         if filename is not None:
-            if not isinstance(filename, str):
+            if prefix is not None or suffix is not None:
                 raise lena.core.LenaTypeError(
-                    arg_error.format("filename", filename)
+                    "filename is incompatible with "
+                    "prefix and suffix. Provide them separately"
                 )
-            self._filename = lena.context.format_context(filename)
+
+        # since the ordering of options is not important,
+        # methods could be a dict.
+        methods = []
+        if prefix is not None:
+            methods.append(("prefix", lena.context.format_context(prefix)))
+        if suffix is not None:
+            methods.append(("suffix", lena.context.format_context(suffix)))
+
+        if filename is not None:
+            methods.append(("filename", lena.context.format_context(filename)))
         if dirname is not None:
-            if not isinstance(dirname, str):
-                raise lena.core.LenaTypeError(
-                    arg_error.format("dirname", dirname)
-                )
-            self._dirname = lena.context.format_context(dirname)
+            methods.append(("dirname", lena.context.format_context(dirname)))
         if fileext is not None:
-            if not isinstance(fileext, str):
-                raise lena.core.LenaTypeError(
-                    arg_error.format("fileext", fileext)
-                )
-            self._fileext = lena.context.format_context(fileext)
+            methods.append(("fileext", lena.context.format_context(fileext)))
+
+        if not methods:
+            # for wrong initialization it must be a TypeError,
+            # as if it was an obligatory argument.
+            raise lena.core.LenaTypeError(
+                "MakeFilename must be initialized with at least "
+                "one of {}".format(args_names)
+            )
+
+        self._methods = methods
 
     def __call__(self, value):
         """Add *output* keys to the *value*'s context.
 
         *filename*, *dirname*, *fileext*, if initialized,
-        set respectively *context.output.{filename,dirname,fileext}*.
-        Only those values are transformed
-        that have no corresponding keys
-        (*filename*, *fileext* or *dirname*) in *context.output*
-        and for which the current context can be formatted
-        (contains all necessary keys for any of the format strings).
+        set respectively *context.output.{filename,dirname,fileext}*
+        (if they didn't exist).
+
+        If this elements sets file name
+        and if context contains *output.prefix* or *output.suffix*,
+        they are prepended to or appended after the file name.
+        After that they are removed from *context.output*.
+
+        If this element adds a prefix or a suffix
+        and they exist in the context,
+        then *prefix* is prepended before the existing prefix,
+        and *suffix* is appended after the existing suffix,
+        unless *overwrite* is set to ``True``:
+        in that case they are overwritten.
+        *prefix* and *suffix* always update their existing keys
+        in the context if they could be formatted
+        (which is different for attributes like *filename*).
+
+        If current context can't be formatted
+        (doesn't contain all necessary keys for the format string),
+        a key is not updated.
         """
         context = lena.flow.get_context(value)
         modified = False
 
-        for key in ["filename", "fileext", "dirname"]:
-            if "output" in context and key in context["output"]:
-                if not self._overwrite:
-                    continue
-            meth = getattr(self, "_" + key, None)
-            if meth is not None:
-                try:
-                    res = meth(context)
-                except lena.core.LenaKeyError:
-                    continue
-                else:
-                    update = {"output": {key: res}}
-                    lena.context.update_recursively(context, update)
-                    modified = True
+        for key, meth in self._methods:
+            if key in ["filename", "fileext", "dirname"]:
+                if "output" in context and key in context["output"]:
+                    if not self._overwrite:
+                        continue
+            try:
+                res = meth(context)
+            except lena.core.LenaKeyError:
+                continue
+            else:
+                if key in ["prefix", "suffix"]:
+                    existing = lena.context.get_recursively(
+                        context, "output." + key, None
+                    )
+                    if existing and not self._overwrite:
+                        # prefix is prepended to existing,
+                        # suffix is appended after existing
+                        if key == "prefix":
+                            res = res + existing
+                        else:
+                            res = existing + res
+                elif key == "filename":
+                    # append existing prefixes and suffixes to filename
+                    prefix = lena.context.get_recursively(
+                        context, "output.prefix", ""
+                    )
+                    suffix = lena.context.get_recursively(
+                        context, "output.suffix", ""
+                    )
+                    res = prefix + res + suffix
+                    # todo: do I need an option *remove_used*?
+                    # I think not. This deletion is natural,
+                    # used parts are conserved in created filename.
+                    # If filename was not created,
+                    # they remain in the context.
+                    if prefix:
+                        del context["output"]["prefix"]
+                    if suffix:
+                        del context["output"]["suffix"]
+                update = {"output": {key: res}}
+                lena.context.update_recursively(context, update)
+                modified = True
 
+        # todo: probably write somewhere how good it was and delete.
+        # just get_data_context in the beginning
+        # and don't check for modifications.
         if modified:
             # or created
             data = lena.flow.get_data(value)

lena/output/write_root_tree.py

@@ -1,3 +1,4 @@
+# needs ROOT installed
 from __future__ import print_function
 
 import array

tests/output/test_make_filename.py

@@ -4,6 +4,8 @@
 import pytest
 
 import lena.core, lena.output
+from lena.core import Sequence
+from lena.flow import Print
 from lena.output import MakeFilename
 from lena.context import format_context
 
@@ -12,12 +14,6 @@ def test_make_filename_init():
     # all format arguments must be strings
     with pytest.raises(lena.core.LenaTypeError):
         MakeFilename(0)
-    with pytest.raises(lena.core.LenaTypeError):
-        MakeFilename(filename=0)
-    with pytest.raises(lena.core.LenaTypeError):
-        MakeFilename(dirname=0)
-    with pytest.raises(lena.core.LenaTypeError):
-        MakeFilename(fileext=0)
 
     # zero arguments are prohibited
     with pytest.raises(lena.core.LenaTypeError):
@@ -27,6 +23,11 @@ def test_make_filename_init():
     with pytest.raises(lena.core.LenaValueError):
         MakeFilename(filename="{}}")
 
+    # prefix and suffix must be provided separately with filename
+    with pytest.raises(lena.core.LenaTypeError):
+        MakeFilename(filename="filename", prefix="wrong_")
+
+
 def test_make_filename():
     data = [(0, {"output": {"filename": "exists"}}),
             (1, {"output": {"fileext": "ext"}}),
@@ -85,3 +86,47 @@ def test_make_filename():
     # context is unchanged
     d = {}
     assert mk((0, d))[1] is d
+
+
+def test_make_filename_prefix_suffix():
+    data = [(0, {"output": {"filename": "file_name"}})]
+    mk1 = MakeFilename(suffix="_suff1", prefix="pref1_")
+    mk2 = MakeFilename(suffix="_suff2", prefix="pref2_")
+
+    # existing file names are unchanged
+    seq12_1 = Sequence(mk1, Print(), mk2)
+    assert list(seq12_1.run(copy.deepcopy(data))) == [(
+        0, {'output':
+            {'filename': 'file_name',
+             'prefix': 'pref2_pref1_',
+             'suffix': '_suff1_suff2'}}
+    )]
+
+    # prefix and suffix are properly added
+    seq12_2 = Sequence(mk1, mk2, MakeFilename("filename"))
+    assert list(seq12_2.run([0])) == [(
+        0, {'output': {'filename': 'pref2_pref1_filename_suff1_suff2'}}
+    )]
+
+    ## Overwrite works
+    # it doesn't harm when file name is produced
+    mk3 = MakeFilename(suffix="_suff2", overwrite=True)
+    seq13 = Sequence(mk1, mk3, MakeFilename("filename3"))
+    assert list(seq13.run(copy.deepcopy(data))) == [
+        (0, {'output':
+             {'prefix': 'pref1_', 'suffix': '_suff2', 
+              'filename': 'file_name'}})
+    ]
+
+    # it really works
+    seq13_2 = Sequence(mk1, mk3, MakeFilename("filename132"))
+    assert list(seq13_2.run([0])) == [
+        (0, {'output': {'filename': 'pref1_filename132_suff2'}})
+    ]
+
+    # formatting arguments work
+    mk4 = MakeFilename(prefix="{{a}}_")
+    seq14 = Sequence(mk1, mk4, MakeFilename("filename14"))
+    assert list(seq14.run([(0, {"a": "A"})])) == [
+        (0, {'a': 'A', 'output': {'filename': 'A_pref1_filename14_suff1'}})
+    ]