Allow get_state_json to accept direct URLs (#271)

* add type hinting to jsonservice

* add get_state test

* omit linting of string type hint

caveclient/jsonservice.py

@@ -2,7 +2,7 @@
 import numbers
 import os
 import re
-from typing import Optional
+from typing import Literal, Optional, Union
 
 import numpy as np
 
@@ -21,6 +21,8 @@
 )
 
 SERVER_KEY = "json_server_address"
+SEUNGLAB_NAMES = ["seunglab"]
+SPELUNKER_NAMES = ["spelunker", "mainline", "cave-explorer"]
 
 
 def neuroglancer_json_encoder(obj):
@@ -100,20 +102,20 @@ def __init__(
         self._ngl_url = ngl_url
 
     @property
-    def state_service_endpoint(self):
+    def state_service_endpoint(self) -> str:
         """Endpoint URL for posting JSON state"""
         url_mapping = self.default_url_mapping
         return self._endpoints["upload_state"].format_map(url_mapping)
 
     @property
-    def ngl_url(self):
+    def ngl_url(self) -> str:
         return self._ngl_url
 
     @ngl_url.setter
     def ngl_url(self, new_ngl_url):
         self._ngl_url = new_ngl_url
 
-    def get_neuroglancer_info(self, ngl_url=None):
+    def get_neuroglancer_info(self, ngl_url: Optional[str] = None) -> dict:
         """Get the info field from a Neuroglancer deployment
 
         Parameters
@@ -142,28 +144,36 @@ def get_neuroglancer_info(self, ngl_url=None):
         handle_response(response, as_json=False)
         return json.loads(response.content)
 
-    def get_state_json(self, state_id):
+    def get_state_json(
+        self,
+        state_id: Union[int, str],
+    ) -> dict:
         """Download a Neuroglancer JSON state
 
         Parameters
         ----------
-        state_id : int
-            ID of a JSON state uploaded to the state service.
+        state_id : int or string
+            ID of a JSON state uploaded to the state service. If a string, treats this as the url to the json state.
 
         Returns
         -------
         dict
             JSON specifying a Neuroglancer state.
         """
-        url_mapping = self.default_url_mapping
-        url_mapping["state_id"] = state_id
-        url = self._endpoints["get_state"].format_map(url_mapping)
+        url = None
+        try:
+            state_id = int(state_id)
+            url_mapping = self.default_url_mapping
+            url_mapping["state_id"] = state_id
+            url = self._endpoints["get_state"].format_map(url_mapping)
+        except ValueError:
+            url = state_id
         response = self.session.get(url)
         handle_response(response, as_json=False)
         return json.loads(response.content)
 
     @_check_version_compatibility(method_constraint=">=0.4.0")
-    def get_property_json(self, state_id):
+    def get_property_json(self, state_id: int) -> dict:
         """Download a Neuroglancer JSON state
 
         Parameters
@@ -183,7 +193,12 @@ def get_property_json(self, state_id):
         handle_response(response, as_json=False)
         return json.loads(response.content)
 
-    def upload_state_json(self, json_state, state_id=None, timestamp=None):
+    def upload_state_json(
+        self,
+        json_state: dict,
+        state_id: Optional[int] = None,
+        timestamp: Optional["time.time"] = None,  # noqa: F821
+    ) -> int:
         """Upload a Neuroglancer JSON state
 
         Parameters
@@ -224,10 +239,10 @@ def upload_state_json(self, json_state, state_id=None, timestamp=None):
     @_check_version_compatibility(">=0.4.0")
     def upload_property_json(
         self,
-        property_json,
-        state_id=None,
+        property_json: dict,
+        state_id: Optional[int] = None,
         max_size: Optional[int] = 2_500_000,
-    ):
+    ) -> int:
         """Upload a Neuroglancer JSON state
 
         Parameters
@@ -277,7 +292,12 @@ def upload_property_json(
         response_re = re.search(".*\/(\d+)", str(response.content))
         return int(response_re.groups()[0])
 
-    def save_state_json_local(self, json_state, filename, overwrite=False):
+    def save_state_json_local(
+        self,
+        json_state: dict,
+        filename: str,
+        overwrite: bool = False,
+    ) -> None:
         """Save a Neuroglancer JSON state to a JSON file locally.
 
         Parameters
@@ -300,12 +320,14 @@ def save_state_json_local(self, json_state, filename, overwrite=False):
 
     def build_neuroglancer_url(
         self,
-        state_id,
-        ngl_url=None,
-        target_site=None,
-        static_url=False,
-        format_properties=False,
-    ):
+        state_id: int,
+        ngl_url: Optional[str] = None,
+        target_site: Optional[
+            Literal["seunglab", "cave-explorer", "mainline", "spelunker"]
+        ] = None,
+        static_url: bool = False,
+        format_properties: bool = False,
+    ) -> str:
         """Build a URL for a Neuroglancer deployment that will automatically retrieve specified state.
         If the datastack is specified, this is prepopulated from the info file field "viewer_site".
         If no ngl_url is specified in either the function or the client, a fallback neuroglancer deployment is used.
@@ -339,17 +361,17 @@ def build_neuroglancer_url(
         if target_site is None and ngl_url is not None:
             ngl_info = self.get_neuroglancer_info(ngl_url)
             if len(ngl_info) > 0:
-                target_site = "cave-explorer"
+                target_site = SPELUNKER_NAMES[0]
             else:
-                target_site = "seunglab"
+                target_site = SEUNGLAB_NAMES[0]
 
         if target_site == "seunglab":
             if ngl_url[-1] == "/":
                 parameter_text = "?json_url="
             else:
                 parameter_text = "/?json_url="
             auth_text = ""
-        elif target_site == "cave-explorer" or target_site == "mainline":
+        elif target_site in SPELUNKER_NAMES:
             if ngl_url[-1] == "/":
                 parameter_text = "#!"
             else:

tests/test_jsonclient.py

@@ -0,0 +1,23 @@
+import responses
+
+from .conftest import datastack_dict
+
+
+@responses.activate()
+def test_basic_state(myclient):
+    global_server = datastack_dict["global_server"]
+    state_id = 1234
+    url = f"{global_server}/nglstate/api/v1/{state_id}"
+    responses.add(responses.GET, url=url, json={"layers": ["img", "seg"]}, status=200)
+    state = myclient.state.get_state_json(state_id)
+    assert "img" in state["layers"]
+
+    direct_url = "https://my-fake-url.com/direct_state.json"
+    responses.add(
+        responses.GET,
+        url=direct_url,
+        json={"layers": ["direct_img", "direct_seg"]},
+        status=200,
+    )
+    state = myclient.state.get_state_json(direct_url)
+    assert "direct_seg" in state["layers"]