From 5515add44d4c948110e995596a66b00a0ac32809 Mon Sep 17 00:00:00 2001
From: Bruno Reis <85531625+brunocbreis@users.noreply.github.com>
Date: Tue, 21 Feb 2023 18:42:20 -0300
Subject: [PATCH] add comp.__contains__ method

fix checks before self.tools was initialized
---
 CHANGELOG.md              |  8 ++++++++
 TODO.md                   |  1 +
 pysion/composition.py     | 25 ++++++++++++++++++++-----
 setup.cfg                 |  2 +-
 tests/test_composition.py | 33 +++++++++++++++++++++++++++++++++
 5 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 70fe9e2..8da442a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,14 @@
 
 ## Unreleased
 
+### Added
+
+- new comp \__contains__ method allows for quick checking if tool in comp even when self.tools is None
+
+### Fixed
+
+- add_merge doesn't break when comp is empty
+
 ## v0.1.1 – 2023-02-21
 
 Some bug fixes and minor changes and improvements.
diff --git a/TODO.md b/TODO.md
index 15f61a2..ceb1556 100644
--- a/TODO.md
+++ b/TODO.md
@@ -14,6 +14,7 @@
   - [ ] basic_comp() for jumpstarting projects
   - [ ] default_comp() that uses configurable settings?
   - [ ] quick titles or something
+- [X] implement \__contains__ method to avoid bugs when self.tools haven't been initialized
 
 ## Docs
 
diff --git a/pysion/composition.py b/pysion/composition.py
index f1bbc95..e88f57e 100644
--- a/pysion/composition.py
+++ b/pysion/composition.py
@@ -114,6 +114,22 @@ def __setitem__(self, key: str, value: Tool | Macro):
 
         self.add_tools(value)
 
+    def __contains__(self, value: Tool | Macro | Modifier | BezierSpline) -> bool:
+        if not (self.tools or self.modifiers):
+            return False
+
+        match value:
+            case Modifier() | BezierSpline():
+                if not self.modifiers:
+                    return False
+                return value in self.modifiers.values()
+            case Tool() | Macro():
+                if not self.tools:
+                    return False
+                return value in self.tools.values()
+            case _:
+                raise ValueError
+
     # Private methods
     def _auto_set_active_tool(self) -> None:
         if not self.tools:
@@ -278,16 +294,15 @@ def add_merge(
             )
 
         # Add tools to comp if not already
-        if background not in self.tools.values():
+        if background not in self:
             if background is not None:
                 self.add_tools(background)
 
-        if foreground not in self.tools.values():
+        if foreground not in self:
             if foreground is not None:
                 self.add_tools(foreground)
 
         merge.add_inputs(bg_input, fg_input)
-
         self.add_tools(merge)
 
         return merge
@@ -363,7 +378,7 @@ def animate(
         """
         match tool:
             case Tool():
-                if tool not in self.tools.values():
+                if tool not in self:
                     print(f"Adding {tool.name} to the comp.\n")
                     self.add_tools(tool)
 
@@ -420,7 +435,7 @@ def animate_position(
         """
         match tool:
             case Tool():
-                if tool not in self.tools.values():
+                if tool not in self:
                     print(f"Adding {tool.name} to the comp.\n")
                     self.add_tools(tool)
 
diff --git a/setup.cfg b/setup.cfg
index 696548d..7b7733a 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = pysion
-version = 0.1.1
+version = 0.1.2
 author = Bruno Reis
 description = Package for creating code that is readable by DaVinci Resolve Fusion.
 url = https://github.com/brunocbreis/pysion
diff --git a/tests/test_composition.py b/tests/test_composition.py
index 4711dd3..af8d664 100644
--- a/tests/test_composition.py
+++ b/tests/test_composition.py
@@ -104,3 +104,36 @@ def test_merge_macro(comp: Composition):
     comp.add_merge("MergeMacro", macro, comp["MyBackground"])
 
     assert comp["MergeMacro"]["Background"].source_operator == macro.name
+
+
+def test_merge_tools_not_in_comp_yet():
+    comp = Composition()
+    bg1 = Tool(ToolID.background, "BG1")
+    bg2 = Tool(ToolID.background, "BG2")
+
+    merge = comp.add_merge("MergeBGs", bg1, bg2)
+
+    assert bg1 in comp.tools.values()
+    assert bg2 in comp.tools.values()
+    assert merge in comp.tools.values()
+
+
+def test_composition_contains_tool(comp: Composition):
+    assert comp["MyBackground"] in comp
+
+
+def test_composition_contains_modifier(comp: Composition):
+    pub = comp.publish(comp["MyBackground"], "TopLeftRed", 1)
+
+    assert pub in comp
+
+
+def test_composition_does_not_contain(comp):
+    new_tool = Tool(ToolID.channel_booleans, "random")
+
+    assert not new_tool in comp
+
+
+def test_composition_contains_error(comp):
+    with pytest.raises(ValueError):
+        "Pizza" in comp