From e38066f6a2c7d4a54fd6c070dd29bd1831c37053 Mon Sep 17 00:00:00 2001
From: David Poznik <dpoznik@23andme.com>
Date: Tue, 19 Mar 2024 16:53:12 -0700
Subject: [PATCH] Reset `Node.hg_snp_set` when `Tree` is reinstantiated (#30)

---
 CHANGELOG.md       |  8 ++++++++
 tests/test_tree.py | 17 +++++++++++++++++
 yhaplo/node.py     | 21 ++++++++++-----------
 3 files changed, 35 insertions(+), 11 deletions(-)
 create mode 100644 tests/test_tree.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6166963..a2f4506 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,14 @@
 Format based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 
+## [2.1.9]
+
+### Fixed
+- `Node.hg_snp` values would drift upon repeated `Tree` instantiation
+
+[2.1.9]: https://github.com/23andMe/yhaplo/compare/2.1.8..2.1.9
+
+
 ## [2.1.8]
 
 ### Added
diff --git a/tests/test_tree.py b/tests/test_tree.py
new file mode 100644
index 0000000..374cf8a
--- /dev/null
+++ b/tests/test_tree.py
@@ -0,0 +1,17 @@
+import numpy as np
+
+from yhaplo.config import Config
+from yhaplo.tree import Tree
+from yhaplo.utils.context_managers import logging_disabled
+
+
+def test_hg_snp_idempotency():
+    config = Config(suppress_output=True)
+    with logging_disabled():
+        tree_1 = Tree(config)
+        tree_2 = Tree(config)
+
+    hg_snps_1 = np.array([node.hg_snp for node in tree_1.depth_first_node_list])
+    hg_snps_2 = np.array([node.hg_snp for node in tree_2.depth_first_node_list])
+
+    assert (hg_snps_1 == hg_snps_2).all()
diff --git a/yhaplo/node.py b/yhaplo/node.py
index 184dc8d..98595e6 100644
--- a/yhaplo/node.py
+++ b/yhaplo/node.py
@@ -56,7 +56,7 @@ class Node:
     tree: "tree_module.Tree"
     config: Config
     args: argparse.Namespace
-    hg_snp_set: set[str] = set()
+    hg_snp_set: set[str]
 
     def __init__(
         self,
@@ -77,7 +77,7 @@ def __init__(
         self.parent = parent
         if parent is None:
             if tree is not None:
-                type(self).set_tree_config_and_args(tree)
+                type(self).set_class_variables(tree)
             else:
                 raise ValueError("Root node requires a tree instance")
 
@@ -204,12 +204,13 @@ def most_highly_ranked_dropped_marker(self) -> "snp_module.DroppedMarker":
     # Class methods
     # ----------------------------------------------------------------------
     @classmethod
-    def set_tree_config_and_args(cls, tree: "tree_module.Tree") -> None:
+    def set_class_variables(cls, tree: "tree_module.Tree") -> None:
         """Set tree, config, and args."""
 
         cls.tree = tree
         cls.config = tree.config
         cls.args = tree.args
+        cls.hg_snp_set = set()
 
     @classmethod
     def truncate_haplogroup_label(cls, haplogroup: str) -> str:
@@ -308,14 +309,12 @@ def priority_sort_snp_list_and_set_hg_snp(self) -> None:
                 self.hg_snp = self.parent.hg_snp + symbol
 
                 # Uniquify if necessary
-                if self.hg_snp in type(self).hg_snp_set:
-                    i = 1
-                    hg_snp_uniqe = f"{self.hg_snp}{i}"
-                    while hg_snp_uniqe in type(self).hg_snp_set:
-                        i += 1
-                        hg_snp_uniqe = f"{self.hg_snp}{i}"
-
-                    self.hg_snp = hg_snp_uniqe
+                original_hg_snp = self.hg_snp
+                i = 0
+                while self.hg_snp in type(self).hg_snp_set:
+                    i += 1
+                    self.hg_snp = f"{original_hg_snp}{i}"
+
             else:
                 logger.warning(
                     "WARNING. Attempted to set star label, "