diff --git a/hbutils/string/template.py b/hbutils/string/template.py index ad38b5ddc85..18cdfab131f 100644 --- a/hbutils/string/template.py +++ b/hbutils/string/template.py @@ -3,12 +3,19 @@ Useful utilities for template a string. """ from string import Template -from typing import Optional, Mapping +from typing import Optional, Mapping, Any -__all__ = ['env_template'] +from ..design import SingletonMark +__all__ = [ + 'env_template' +] -def env_template(template: str, environ: Optional[Mapping[str, str]] = None, safe: bool = False) -> str: +_NO_DEFAULT_VALUE = SingletonMark('_NO_DEFAULT_VALUE') + + +def env_template(template: str, environ: Optional[Mapping[str, Any]] = None, + safe: bool = False, default: Any = _NO_DEFAULT_VALUE) -> str: """ Overview: Mapping all the environment values (not system environment variables) into a template string. @@ -18,6 +25,8 @@ def env_template(template: str, environ: Optional[Mapping[str, str]] = None, saf - environ (:obj:`Optional[Mapping[str, str]]`): Environment values, should be a mapping. - safe (:obj:`bool`): Safe substitute, default is ``False`` which means \ all the value used in template should be able to found in the given ``environ``. + - default (:obj:`Any`): Default value when no variable provided. Default is ``_NO_DEFAULT_VALUE``, which \ + means ``KeyError`` will be raised. Returns: - result (:obj:`str`): Substituted string. @@ -30,7 +39,19 @@ def env_template(template: str, environ: Optional[Mapping[str, str]] = None, saf KeyError: 'B' >>> env_template('${A} + 1 = ${B}', {'A': '1'}, safe=True) '1 + 1 = ${B}' + >>> env_template('${A} + 1 = ${B}', {'A': '1'}, default='') # like environment variable + '1 + 1 = ' """ + + class _DefaultDict(dict): + def __getitem__(self, item): + return dict.get(self, item, default) + _template = Template(template) + env = environ or {} + if default is not _NO_DEFAULT_VALUE: + env = _DefaultDict(env) + _func = _template.safe_substitute if safe else _template.substitute - return _func(**(environ or {})) + # noinspection PyArgumentList + return _func(env) diff --git a/test/string/test_template.py b/test/string/test_template.py index d174b026580..1e1d3de9b8e 100644 --- a/test/string/test_template.py +++ b/test/string/test_template.py @@ -10,3 +10,7 @@ def test_env_template(self): assert env_template('this is the ${number} day.', dict(number='1st')) == 'this is the 1st day.' with pytest.raises(KeyError): env_template('this is the ${number} day') + + def test_env_template_with_default(self): + assert env_template('${A} + 1 == ${B}', {'A': 1}, default='') == '1 + 1 == ' + assert env_template('${A} + 1 == ${B}', {'A': 1}, default='', safe=True) == '1 + 1 == '