From 808143cc2543d04a844bbaefd34b0e2d15c8b977 Mon Sep 17 00:00:00 2001
From: Nao Pross <np@0hm.ch>
Date: Sun, 10 Mar 2024 23:42:16 +0100
Subject: Fix PolyIndex.__contains__

---
 mdpoly/index.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/mdpoly/index.py b/mdpoly/index.py
index 5ff28b6..5211245 100644
--- a/mdpoly/index.py
+++ b/mdpoly/index.py
@@ -160,8 +160,13 @@ class PolyIndex(tuple[PolyVarIndex]):
         indices = ", ".join(map(repr, self))
         return f"{name}({indices})"
 
-    def __contains__(self, var_idx):
-        return any(map(lambda e: e.var_idx == var_idx, self))
+    # FIXME: may need to convert to method to avoid breaking expected behaviour
+    def __contains__(self, index: PolyVarIndex | int) -> bool:
+        if isinstance(index, PolyVarIndex):
+            return tuple.__contains__(self, index)
+
+        if isinstance(index, int):
+            return any(map(lambda e: e.var_idx == index, self))
 
     # -- Helper methods ---
 
-- 
cgit v1.2.1