Add less comparator for PeerID (#353)

* add less for PeerID * Support non-PeerID types, use .to_base58() instead of ._bytes since only it is visible to humans * Use TypeError instead of ValueError * Fix __less__ -> __lt__, add tests Co-authored-by: Alexander Borzunov <[email protected]>
learning-at-home · Aug 18, 2021 · dd05abb · dd05abb
1 parent cf31951
commit dd05abb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/hivemind/p2p/p2p_daemon_bindings/datastructures.py b/hivemind/p2p/p2p_daemon_bindings/datastructures.py
@@ -74,6 +74,12 @@ def __eq__(self, other: object) -> bool:
         else:
             return False
 
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, PeerID):
+            raise TypeError(f"'<' not supported between instances of 'PeerID' and '{type(other)}'")
+
+        return self.to_base58() < other.to_base58()
+
     def __hash__(self) -> int:
         return hash(self._bytes)
 

diff --git a/tests/test_p2p_daemon_bindings.py b/tests/test_p2p_daemon_bindings.py
@@ -144,6 +144,12 @@ def test_peer_id():
     peer_id_3 = PeerID.from_base58("QmbmfNDEth7Ucvjuxiw3SP3E4PoJzbk7g4Ge6ZDigbCsNp")
     assert PEER_ID != peer_id_3
 
+    a = PeerID.from_base58("bob")
+    b = PeerID.from_base58("eve")
+    assert a < b and b > a and not (b < a) and not (a > b)
+    with pytest.raises(TypeError):
+        assert a < object()
+
 
 def test_stream_info():
     proto = "123"