fix #111 add timeout option to crossref class and its methods

sckott · Nov 8, 2024 · 6a2005f · 6a2005f
1 parent 99f1ab6
commit 6a2005f
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 37 deletions.
diff --git a/habanero/__init__.py b/habanero/__init__.py
@@ -37,7 +37,7 @@
 """
 
 __title__ = "habanero"
-__version__ = "1.2.6.91"
+__version__ = "1.2.7.91"
 __author__ = "Scott Chamberlain"
 __license__ = "MIT"
 

diff --git a/habanero/crossref/crossref.py b/habanero/crossref/crossref.py
@@ -18,6 +18,7 @@ class Crossref:
     :param api_key: An API key to send with each http request
     :param mailto: A mailto string, see section below
     :param ua_string: A user agent string, see section below
+    :param timeout: curl timeout
 
     |
     |
@@ -169,21 +170,24 @@ def __init__(
         api_key: str = None,
         mailto: str = None,
         ua_string: str = None,
+        timeout: int = 5,
     ) -> None:
         self.base_url = base_url
         self.api_key = api_key
         self.mailto = mailto
         self.ua_string = ua_string
+        self.timeout = timeout
 
     def __repr__(self):
         return (
-            """< %s \nURL: %s\nKEY: %s\nMAILTO: %s\nADDITIONAL UA STRING: %s\n>"""
+            """< %s \nURL: %s\nKEY: %s\nMAILTO: %s\nADDITIONAL UA STRING: %s\nTimeout: %s\n>"""
             % (
                 type(self).__name__,
                 self.base_url,
                 sub_str(self.api_key),
                 self.mailto,
                 self.ua_string,
+                self.timeout,
             )
         )
 
@@ -337,9 +341,7 @@ def works(
         """
         if ids.__class__.__name__ != "NoneType":
             return request(
-                self.mailto,
-                self.ua_string,
-                self.base_url,
+                self,
                 "/works/",
                 ids,
                 query,
@@ -363,6 +365,7 @@ def works(
             return Request(
                 self.mailto,
                 self.ua_string,
+                self.timeout,
                 self.base_url,
                 "/works/",
                 query,
@@ -476,9 +479,7 @@ def members(
             res = cr.members(filter = {'has_public_references': True})
         """
         return request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/members/",
             ids,
             query,
@@ -594,9 +595,7 @@ def prefixes(
         """
         check_kwargs(["query"], kwargs)
         return request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/prefixes/",
             ids,
             query=None,
@@ -723,9 +722,7 @@ def funders(
             [type(w) for w in x] # [dict, NoneType]
         """
         return request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/funders/",
             ids,
             query,
@@ -848,9 +845,7 @@ def journals(
             [ x.get('title') for x in res['message']['items'] ]
         """
         return request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/journals/",
             ids,
             query,
@@ -952,9 +947,7 @@ def types(
             [ x.get('title') for x in res['message']['items'] ]
         """
         return request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/types/",
             ids,
             query,
@@ -1011,9 +1004,7 @@ def licenses(
         """
         check_kwargs(["ids", "filter", "works"], kwargs)
         res = request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/licenses/",
             None,
             query,
@@ -1063,9 +1054,7 @@ def registration_agency(self, ids: Union[List[str], str], **kwargs) -> list:
             kwargs,
         )
         res = request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/works/",
             ids,
             None,
@@ -1109,9 +1098,7 @@ def random_dois(self, sample: int = 10, **kwargs) -> list:
             cr.random_dois(100)
         """
         res = request(
-            self.mailto,
-            self.ua_string,
-            self.base_url,
+            self,
             "/works/",
             None,
             None,

diff --git a/habanero/request.py b/habanero/request.py
@@ -17,9 +17,7 @@
 
 
 def request(
-    mailto,
-    ua_string,
-    url,
+    cr,
     path,
     ids=None,
     query=None,
@@ -41,7 +39,7 @@ def request(
 ):
     """HTTP request helper."""
     warning_thrown = False
-    url = url + path
+    url = cr.base_url + path
 
     if cursor_max.__class__.__name__ != "NoneType":
         if not isinstance(cursor_max, int):
@@ -77,7 +75,12 @@ def request(
     if ids.__class__.__name__ == "NoneType":
         url = url.strip("/")
         try:
-            r = httpx.get(url, params=payload, headers=make_ua(mailto, ua_string))
+            r = httpx.get(
+                url,
+                params=payload,
+                headers=make_ua(cr.mailto, cr.ua_string),
+                timeout=cr.timeout,
+            )
             r.raise_for_status()
         except httpx.HTTPStatusError:
             if is_json(r):
@@ -103,8 +106,9 @@ def request(
         for i in range(len(ids)):
             if works:
                 res = Request(
-                    mailto,
-                    ua_string,
+                    cr.mailto,
+                    cr.ua_string,
+                    cr.timeout,
                     url,
                     str(ids[i]) + "/works",
                     query,
@@ -131,7 +135,12 @@ def request(
 
                 endpt = endpt.strip("/")
 
-                r = httpx.get(endpt, params=payload, headers=make_ua(mailto, ua_string))
+                r = httpx.get(
+                    endpt,
+                    params=payload,
+                    headers=make_ua(cr.mailto, cr.ua_string),
+                    timeout=cr.timeout,
+                )
                 if r.status_code > 201 and should_warn:
                     warning_thrown = True
                     mssg = "%s on %s: %s" % (r.status_code, ids[i], r.reason_phrase)

diff --git a/habanero/request_class.py b/habanero/request_class.py
@@ -27,6 +27,7 @@ def __init__(
         self,
         mailto,
         ua_string,
+        timeout,
         url,
         path,
         query=None,
@@ -46,6 +47,7 @@ def __init__(
     ):
         self.mailto = mailto
         self.ua_string = ua_string
+        self.timeout = timeout
         self.url = url
         self.path = path
         self.query = query
@@ -148,6 +150,7 @@ def _req(self, payload, should_warn):
                 self._url(),
                 params=payload,
                 headers=make_ua(self.mailto, self.ua_string),
+                timeout=self.timeout,
             )
             r.raise_for_status()
         except httpx.HTTPStatusError: