Merge pull request #4 from corneliusroemer/partial_aliasing

feat: add partial compression with up_to and accepted_aliases

README.md

@@ -26,6 +26,14 @@ aliasor.uncompress("XA") # 'XA'
 
 # Go from unaliased lineage to aliased lineage
 aliasor.compress("B.1.1.529.3.1") # 'BA.3.1'
+
+aliasor.partial_compress("B.1.1.529.3.1",up_to=1) # 'BA.3.1'
+aliasor.partial_compress("B.1.1.529.3.1.1.2",up_to=1) # 'BA.3.1.1.2'
+
+aliasor.partial_compress("B.1.1.529.3.1",accepted_aliases=["AY"]) # 'B.1.1.529.3.1'
+aliasor.partial_compress("B.1.617.2",accepted_aliases=["AY"]) # 'AY.2'
+
+aliasor.partial_compress('B.1.1.529.2.75.1.2',up_to=4, accepted_aliases={"BA"}) == 'BL.2'
 ```
 
 See [tests](tests/test_aliasor.py) for more examples.

src/pango_aliasor/aliasor.py

@@ -47,6 +47,40 @@ def uncompress(self, name):
             return unaliased + "." + name_split[1]
         else:
             return unaliased + "." + ".".join(name_split[1:])
+
+    def partial_compress(self, name, up_to: int = 0, accepted_aliases: set = {}):
+        """
+        aliasor.partial_compress("B.1.1.529.3.1",up_to=1) # 'BA.3.1'
+        aliasor.partial_compress("B.1.1.529.3.1.1.2",up_to=1) # 'BA.3.1.1.2'
+
+        aliasor.partial_compress("B.1.1.529.3.1",accepted_aliases=["AY"]) # 'B.1.1.529.3.1'
+        aliasor.partial_compress("B.1.617.2",accepted_aliases=["AY"]) # 'AY.2'
+        """
+        # If accepted_aliases is passed without up_to set, then try out all possible values
+        name_split = name.split(".")
+        levels = len(name_split) - 1
+        indirections = (levels - 1) // 3
+
+        alias = name_split[0]
+
+        if up_to > 0:
+            if indirections <= up_to:
+                return self.compress(name)
+            to_alias = ".".join(name_split[0 : (3 * up_to + 1)])
+            alias = self.realias_dict[to_alias]
+
+        # Compress at least till up_to, maybe further
+
+        # Check if levels beyond up_to (working backwards) are in accepted_aliases
+        if accepted_aliases is not {}:
+            for level in range(indirections,up_to,-1):
+                to_alias = ".".join(name_split[0 : (3 * level + 1)])
+                if to_alias in self.realias_dict.keys():
+                    if self.realias_dict[to_alias] in accepted_aliases:
+                        alias = self.realias_dict[to_alias]
+                        return alias + "." + ".".join(name_split[(3 * level + 1) :])
+
+        return alias + "." + ".".join(name_split[(3 * up_to + 1) :])
 
 
 # %%

tests/test_aliasor.py

@@ -60,3 +60,23 @@ def test_double_alias_uncompression():
 def test_read_from_file(datadir):
     aliasor = Aliasor(datadir.join('alias_key.json'))
     assert aliasor.compress('B.1.1.529.1') == 'BA.1'
+
+def test_partial_alias_up_to():
+    aliasor = Aliasor()
+    assert aliasor.partial_compress('B.1.1.529.1.2', up_to = 0) == 'B.1.1.529.1.2'
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2', up_to = 1) == 'BA.2.75.1.2'
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2', up_to = 2) == 'BL.2'
+
+def test_partial_alias_accepted():
+    aliasor = Aliasor()
+    assert aliasor.partial_compress('B.1.1.529.1.2', accepted_aliases={"BA","AZ"}) == "BA.1.2"
+    assert aliasor.partial_compress('B.1.617.2.3', accepted_aliases={"BA","AZ"}) == "B.1.617.2.3"
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2', accepted_aliases={"BA"}) == 'BA.2.75.1.2'
+
+def test_partial_alias_combination():
+    aliasor = Aliasor()
+    assert aliasor.partial_compress('B.1.1.529.1.2',up_to=1, accepted_aliases={"BA","AZ"}) == "BA.1.2"
+    assert aliasor.partial_compress('B.1.617.2.3',up_to=1, accepted_aliases={"BA","AZ"}) == "AY.3"
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2',up_to=3, accepted_aliases={"BA"}) == 'BL.2'
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2',up_to=4, accepted_aliases={"BA"}) == 'BL.2'
+    assert aliasor.partial_compress('B.1.1.529.2.75.1.2',up_to=1, accepted_aliases={"BA"}) == 'BA.2.75.1.2'