From f758df7d0cf70c01cd4c0b84ce2c2e3c09e4e8a4 Mon Sep 17 00:00:00 2001
From: Michiel Cottaar <MichielCottaar@protonmail.com>
Date: Thu, 4 Apr 2024 16:35:23 +0100
Subject: [PATCH] Support iteratively linking of placeholder values

---
 CHANGELOG.md               |  2 ++
 src/file_tree/template.py  |  9 ++++++++-
 src/tests/test_template.py | 20 ++++++++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 386e13a..a91a561 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,7 @@
 # Changelog
 ## [Unreleased]
+### Added
+- Linking can now be applied iteratively. For example, after `tree.placeholders.link("A", "B")` and `tree.placeholders.link("B", "C")`, you will now end up with "A", "B", and "C" all linked to each other.
 ## [v1.3.0]
 ### Added
 - A `link` keyword to `update_glob`, which allows one to indicate that two placeholders should co-vary.
diff --git a/src/file_tree/template.py b/src/file_tree/template.py
index 4da929e..189d9dd 100644
--- a/src/file_tree/template.py
+++ b/src/file_tree/template.py
@@ -95,6 +95,7 @@ class Placeholders(MutableMapping):
             value = {k: v for k, v in zip(key, value)}
             key = frozenset(key)
         if isinstance(key, frozenset):
+            assert isinstance(value, dict)
             for k in list(key):
                 if k in self.linkages:
                     unmatched_keys = [
@@ -272,7 +273,13 @@ class Placeholders(MutableMapping):
         This can be thought of using `zip` for linked variables and
         `itertools.product` for unlinked ones.
         """
-        self[keys] = [self[key] for key in keys]
+        actual_keys = set()
+        for key in keys:
+            if key in self.linkages:
+                actual_keys.update(self.linkages[key])
+            else:
+                actual_keys.add(key)
+        self[frozenset(actual_keys)] = {key: self[key] for key in actual_keys}
 
     def unlink(self, *keys):
         """
diff --git a/src/tests/test_template.py b/src/tests/test_template.py
index e5e4bb5..44f0fd4 100644
--- a/src/tests/test_template.py
+++ b/src/tests/test_template.py
@@ -267,6 +267,26 @@ def test_linkages_overwrite():
     assert p["A"] == (1, 1, 2)
     assert p["B"] == (4, 6, 5)
 
+    # test iteratively linking variables
+    p = template.Placeholders()
+    for key in ["A", "B", "C", "D"]:
+        p[key] = (1, 2, 3)
+    assert len(list(p.iter_over(["A", "B", "C", "D"]))) == 3**4
+    p.link("A", "B")
+    p.link("A", "C", "D")
+    assert len(list(p.iter_over(["A", "B", "C", "D"]))) == 3
+
+    p = template.Placeholders()
+    for key in ["A", "B", "C", "D"]:
+        p[key] = (1, 2, 3)
+    assert len(list(p.iter_over(["A", "B", "C", "D"]))) == 3**4
+    p.link("A", "B")
+    p.link("C", "D")
+    assert len(list(p.iter_over(["A", "B", "C", "D"]))) == 3**2
+    p.link("A", "B", "C")
+    assert len(list(p.iter_over(["A", "B", "C", "D"]))) == 3
+
+
 
 def test_missing_keys():
     p = template.Placeholders()
-- 
GitLab