Merge pull request #4 from oughtinc/join

F.Join

fvalues/f.py

@@ -1,6 +1,7 @@
 import ast
 import inspect
 import warnings
+from collections.abc import Iterable
 from copy import deepcopy
 from dataclasses import dataclass
 from functools import lru_cache
@@ -126,8 +127,6 @@ def _parts_from_node(
             f_value = FValue(source, value, formatted)
             return (f_value,)
         else:
-            # Part of a concatenation.
-            assert isinstance(node.parent, (ast.BinOp, ast.AugAssign))  # type: ignore
             assert isinstance(value, str)
             f_value = FValue(get_node_source_text(node, ex.source), value, value)
             return (f_value,)
@@ -263,6 +262,44 @@ def __add__(self, other: str) -> "F":
     def __radd__(self, other: str) -> "F":
         return self._add(other, False)
 
+    def join(self, iterable: Iterable[str]) -> "F":
+        parts: list[Part] = []
+        to_list = not isinstance(iterable, (list, tuple))
+        if to_list:
+            iterable = list(iterable)
+        ex = executing.Source.executing(get_frame())
+        iterable_source = None
+        separator_source = None
+        if (
+            ex.node
+            and isinstance(ex.node, ast.Call)
+            and isinstance(ex.node.func, ast.Attribute)
+            and ex.node.func.attr == "join"
+            and len(ex.node.args) == 1
+        ):
+            [iterable_node] = ex.node.args
+            iterable_source = get_node_source_text(iterable_node, ex.source)
+            iterable_source = f"({iterable_source})"
+            if to_list:
+                iterable_source = f"list{iterable_source}"
+
+            separator_node = ex.node.func.value
+            separator_source = get_node_source_text(separator_node, ex.source)
+
+        for i, item in enumerate(iterable):
+            assert isinstance(item, str)
+            if i:
+                if separator_source:
+                    parts.append(FValue(separator_source, self, str(self)))
+                else:
+                    parts.append(self)
+
+            if iterable_source:
+                parts.append(FValue(f"{iterable_source}[{i}]", item, str(item)))
+            else:
+                parts.append(item)
+        return F(str(self).join(map(str, iterable)), tuple(parts))
+
 
 def get_frame() -> FrameType:
     """

tests/test_f.py

@@ -223,3 +223,64 @@ def test_bad_source_segment():
     assert part.source in ("1 + (2)", "1 + 2")
     assert part.value == 3
     assert part.formatted == "3"
+
+
+def test_other_node_type_call_arg():
+    s = "foo"
+    s = F(F(s))
+    assert s == "foo"
+    (part,) = s.parts
+    assert part == FValue(source="F(s)", value="foo", formatted="foo")
+    assert isinstance(part, FValue)  # for mypy
+    assert (
+        part.value.parts
+        == s.flatten().parts
+        == (FValue(source="s", value="foo", formatted="foo"),)
+    )
+
+
+def test_join_non_list():
+    strings = (x for x in ["a", "b", "c"])
+    s = F(" ").join(strings)
+    assert s == "a b c"
+    assert s.parts == (
+        FValue(source="list(strings)[0]", value="a", formatted="a"),
+        FValue(source='F(" ")', value=" ", formatted=" "),
+        FValue(source="list(strings)[1]", value="b", formatted="b"),
+        FValue(source='F(" ")', value=" ", formatted=" "),
+        FValue(source="list(strings)[2]", value="c", formatted="c"),
+    )
+    assert s.flatten().parts == (
+        FValue(source="list(strings)[0]", value="a", formatted="a"),
+        " ",
+        FValue(source="list(strings)[1]", value="b", formatted="b"),
+        " ",
+        FValue(source="list(strings)[2]", value="c", formatted="c"),
+    )
+
+
+def test_join_list():
+    strings = ["a", "b", "c"]
+    s = F("").join(strings)
+    assert s == "abc"
+    assert s.parts == (
+        FValue(source="(strings)[0]", value="a", formatted="a"),
+        FValue(source='F("")', value="", formatted=""),
+        FValue(source="(strings)[1]", value="b", formatted="b"),
+        FValue(source='F("")', value="", formatted=""),
+        FValue(source="(strings)[2]", value="c", formatted="c"),
+    )
+    assert s.flatten().parts == (
+        FValue(source="(strings)[0]", value="a", formatted="a"),
+        "",
+        FValue(source="(strings)[1]", value="b", formatted="b"),
+        "",
+        FValue(source="(strings)[2]", value="c", formatted="c"),
+    )
+
+
+def test_join_bad_source():
+    strings = ["a", "b", "c"]
+    s = F.join(F(","), strings)
+    assert s == "a,b,c"
+    assert s.parts == s.flatten().parts == ("a", ",", "b", ",", "c")