Add truncate_middle to jinja env (#258)

* Add truncate_middle to jinja env

* apply lint, fix tests

Author: Ktakuya332C

flexeval/core/utils/jinja2_utils.py

@@ -10,9 +10,19 @@ def regex_replace(s: str, pattern: str, replace: str) -> str:
     return re.sub(pattern, replace, s)
 
 
+def truncate_middle(s: str, length: int, filler: str = "...") -> str:
+    if len(s) + len(filler) <= length:
+        return s
+    extract_len = length - len(filler)
+    prefix_len = extract_len // 2
+    suffix_len = extract_len // 2 + extract_len % 2
+    return s[:prefix_len] + filler + s[-suffix_len:]
+
+
 JINJA2_ENV = sandbox.ImmutableSandboxedEnvironment(
     undefined=jinja2.StrictUndefined,
     autoescape=False,  # important for not escaping double quotations
     keep_trailing_newline=True,
 )
 JINJA2_ENV.filters["regex_replace"] = regex_replace
+JINJA2_ENV.filters["truncate_middle"] = truncate_middle

tests/core/utils/test_jinja2_env.py

@@ -6,3 +6,12 @@
 def test_regex_replace() -> None:
     template = "{{ text | regex_replace('<<.*?>>', '') }}"
     assert JINJA2_ENV.from_string(template).render(text="<<a>>Hello <<dummy>>world!") == "Hello world!"
+
+
+def test_truncate_middle() -> None:
+    template = JINJA2_ENV.from_string("{{ text | truncate_middle(14, '+') }}")
+    assert template.render(text="This is a pen") == "This is a pen"
+    template = JINJA2_ENV.from_string("{{ text | truncate_middle(13, '+') }}")
+    assert template.render(text="This is a pen") == "This i+ a pen"
+    template = JINJA2_ENV.from_string("{{ text | truncate_middle(12, '+') }}")
+    assert template.render(text="This is a pen") == "This + a pen"