From 4048fb78f46d8aff4546c8d4e7fad82f230c5c00 Mon Sep 17 00:00:00 2001
From: Hugo Cachitas <hcachitas@gmail.com>
Date: Mon, 1 Apr 2024 00:18:02 +0100
Subject: [PATCH] Add network API endpoint

---
 src/stringx/__init__.py |  5 +++++
 src/stringx/client.py   | 24 ++++++++++++++++++++++
 tests/test_client.py    | 45 +++++++++++++++++++++++++++++++++++++++++
 tests/test_stringx.py   | 22 ++++++++++++++++++++
 4 files changed, 96 insertions(+)

diff --git a/src/stringx/__init__.py b/src/stringx/__init__.py
index 06393b3..0771bac 100644
--- a/src/stringx/__init__.py
+++ b/src/stringx/__init__.py
@@ -18,6 +18,11 @@ def map_identifiers(identifiers: List[str], species: int):
         return client.map_identifiers(identifiers=identifiers, species=species)
 
 
+def network(identifiers: List[str], species: int):
+    with Client() as client:
+        return client.network(identifiers=identifiers, species=species)
+
+
 def interaction_partners(identifiers: List[str], species: int):
     with Client() as client:
         return client.interaction_partners(identifiers=identifiers, species=species)
diff --git a/src/stringx/client.py b/src/stringx/client.py
index a1cb00c..accea18 100644
--- a/src/stringx/client.py
+++ b/src/stringx/client.py
@@ -48,6 +48,30 @@ def map_identifiers(
         }
         return self.request("get_string_ids", params=params)
 
+    def network(
+        self,
+        identifiers: List[str],
+        species: int,
+        required_score: Optional[float] = None,
+        network_type: str = "functional",
+        add_nodes: Optional[int] = None,
+        show_query_node_labels: bool = False,
+    ):
+        params = dict(
+            identifiers="\r".join(identifiers),
+            species=species,
+            network_type=network_type,
+            show_query_node_labels=int(show_query_node_labels),
+        )
+
+        if required_score:
+            params |= {"required_score": required_score}
+
+        if add_nodes:
+            params |= {"add_nodes": add_nodes}
+
+        return self.request("network", params=params)
+
     def interaction_partners(
         self,
         identifiers: List[str],
diff --git a/tests/test_client.py b/tests/test_client.py
index 57135b3..a392906 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -55,6 +55,51 @@ def test_map_identifiers(httpx_mock):
         client.map_identifiers(["some_identifier"], 7227)
 
 
+def test_network(httpx_mock):
+    httpx_mock.add_response(
+        url=httpx.URL(
+            "https://string-db.org/api/json/network",
+            params={
+                "identifiers": "id1",
+                "species": "7227",
+                "network_type": "functional",
+                "show_query_node_labels": 0,
+                "caller_identity": stringx.DEFAULT_CALLER_IDENTITY,
+            },
+        ),
+        method="POST",
+        json=True,
+    )
+
+    httpx_mock.add_response(
+        url=httpx.URL(
+            "https://string-db.org/api/json/network",
+            params={
+                "identifiers": "id1\rid2",
+                "species": "7227",
+                "required_score": 1,
+                "network_type": "physical",
+                "add_nodes": 2,
+                "show_query_node_labels": 1,
+                "caller_identity": stringx.DEFAULT_CALLER_IDENTITY,
+            },
+        ),
+        method="POST",
+        json=True,
+    )
+
+    with stringx.Client() as client:
+        client.network(["id1"], 7227)
+        client.network(
+            identifiers=["id1", "id2"],
+            species=7227,
+            required_score=1,
+            network_type="physical",
+            add_nodes=2,
+            show_query_node_labels=True,
+        )
+
+
 def test_interaction_partners(httpx_mock):
     httpx_mock.add_response(
         url=httpx.URL(
diff --git a/tests/test_stringx.py b/tests/test_stringx.py
index 1cb03c4..f2546c3 100644
--- a/tests/test_stringx.py
+++ b/tests/test_stringx.py
@@ -18,6 +18,28 @@ def test_map_identifiers():
         assert "annotation" in keys
 
 
+def test_network():
+    network = stringx.network(["edin", "atta", "attc"], 7227)
+
+    assert len(network) == 3
+
+    for id_ in network:
+        keys = id_.keys()
+        assert "stringId_A" in keys
+        assert "stringId_B" in keys
+        assert "preferredName_A" in keys
+        assert "preferredName_B" in keys
+        assert "ncbiTaxonId" in keys
+        assert "score" in keys
+        assert "nscore" in keys
+        assert "fscore" in keys
+        assert "pscore" in keys
+        assert "ascore" in keys
+        assert "escore" in keys
+        assert "dscore" in keys
+        assert "tscore" in keys
+
+
 def test_interaction_partners():
     interaction_partners = stringx.interaction_partners(["edin"], 7227)