updated issue with regex parsing the cursormark and updated testing'

Author: digitalsouth

industryDocumentsWrapper/ucsf_api.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 import re
 import requests
-import time
 import polars as pl
 
 
@@ -20,13 +19,13 @@ def _create_query(self, **kwargs) -> str:
         if kwargs['q']:
             query = f"{self.base_url}query?q=({kwargs['q']})&wt={kwargs['wt']}&cursorMark={kwargs['cursorMark']}&sort={kwargs['sort']}"
         else:
-            query = f"{self.base_url}query?q=("+' AND '.join([f'{k}:{v}' for k, v in kwargs.items() if v and k != 'wt' and k != 'cursorMark' and k != 'sort'])+f")&wt={kwargs['wt']}&cursorMark={kwargs['cursorMark']}&sort={kwargs['sort']}"
+            query = f"{self.base_url}query?q=("+' AND '.join([f'{k}:{v}' for k, v in kwargs.items() if v and k != 'wt' and k != 'cursorMark' and k != 'sort' and k != 'n'])+f")&wt={kwargs['wt']}&cursorMark={kwargs['cursorMark']}&sort={kwargs['sort']}"
 
         return query
 
     def _update_cursormark(self, query:str, cursor_mark: str) -> str:
         """Updates cursor mark in query string"""
-        return re.sub(r'(?<=cursorMark=)[A-Za-z0-9*]+(?=&)', cursor_mark, query)
+        return re.sub(r'(?<=cursorMark=)[A-Za-z0-9*=]+(?=&)', cursor_mark, query)
 
     def _loop_results(self, query:str, n:int) -> None:
         """Iteratively retrieves documents with cursor_mark for Solr deep paging"""
@@ -39,12 +38,16 @@ def _loop_results(self, query:str, n:int) -> None:
         while (next_cursor != current_cursor) and (len(self.results) < n):
 
             if next_cursor:
-                query = self._update_cursormark(query, next_cursor)
-                
+                print(current_cursor)
+                current_cursor = next_cursor
+                print(current_cursor)
+                query = self._update_cursormark(query, current_cursor)
+                print(query)
+            
+            print(query)
             r = requests.get(query).json()
 
-            current_cursor = r['responseHeader']['params']['cursorMark']
-
+            print(r['response']['numFound'])
             if n < len(r['response']['docs']):
                 self.results.extend(r['response']['docs'][:n])
 
@@ -53,7 +56,10 @@ def _loop_results(self, query:str, n:int) -> None:
 
             else:
                 self.results.extend(r['response']['docs'])
-                next_cursor = r['nextCursorMark']
+            
+            next_cursor = r['nextCursorMark']
+            
+            print(f'Current cursor: {current_cursor} | Next cursor: {next_cursor}')
 
             print(f"{len(self.results)}/{n} documents collected")
 
@@ -64,7 +70,6 @@ def _create_links(self, industry) -> None:
         for doc in self.results:
             doc['url'] = f"https://www.industrydocuments.ucsf.edu/{industry}/docs/#id={doc['id']}"
 
-    # TODO: add functionality for /select/* queries for specific fields
     def query(self, 
         q:str = False,
         case:str = False,
@@ -110,6 +115,7 @@ def query(self,
 
         """Queries the UCSF Industry Documents Solr Library for documents"""
         self._loop_results(query, n)
+        
         if industry:
             self._create_links(industry)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "industryDocumentsWrapper"
-version = "0.1.2"
+version = "0.1.3"
 description = "A simple python wrapper for the UCSF Industry Documents API."
 authors = ["Rolando Rodriguez <[email protected]>"]
 maintainers = ["Rolando Rodriguez <[email protected]>"]

tests/test_ucsf_api.py

@@ -70,12 +70,14 @@ def test_loop_results_50(indDocSearch):
     query = 'https://metadata.idl.ucsf.edu/solr/ltdl3/query?q=(industry:tobacco AND case:"State of North Carolina" AND collection:"JUUL labs Collection" AND type:email)&wt=json&cursorMark=*&sort=id%20asc'
     indDocSearch._loop_results(query, 50)
     assert len(indDocSearch.results) == 50
+    assert len(set([x['id'] for x in indDocSearch.results])) == 50
     assert indDocSearch.results[0]['id'] == 'ffbb0284'
 
 def test_loop_results_150(indDocSearch):
     query = 'https://metadata.idl.ucsf.edu/solr/ltdl3/query?q=(industry:tobacco AND case:"State of North Carolina" AND collection:"JUUL labs Collection" AND type:email)&wt=json&cursorMark=*&sort=id%20asc'
     indDocSearch._loop_results(query, 150)
     assert len(indDocSearch.results) == 150
+    assert len(set([x['id'] for x in indDocSearch.results])) == 150
     assert indDocSearch.results[0]['id'] == 'ffbb0284'
 
 def test_create_links(indDocSearch, mock_results):
@@ -87,22 +89,26 @@ def test_create_links(indDocSearch, mock_results):
 def test_query_with_q_100(indDocSearch):
     indDocSearch.query(q='industry:tobacco AND case:"State of North Carolina" AND collection:"JUUL labs Collection" AND type:email', n=100)
     assert len(indDocSearch.results) == 100
+    assert len(set([x['id'] for x in indDocSearch.results])) == 100
     assert indDocSearch.results[0]['id'] == 'ffbb0284'
     assert indDocSearch.results[0]['url'] == 'https://www.industrydocuments.ucsf.edu/tobacco/docs/#id=ffbb0284'
 
 def test_query_with_q_500(indDocSearch):
     indDocSearch.query(q='industry:tobacco AND case:"State of North Carolina" AND collection:"JUUL labs Collection" AND type:email', n=500)
     assert len(indDocSearch.results) == 500
+    assert len(set([x['id'] for x in indDocSearch.results])) == 500
     assert indDocSearch.results[0]['id'] == 'ffbb0284'
     assert indDocSearch.results[0]['url'] == 'https://www.industrydocuments.ucsf.edu/tobacco/docs/#id=ffbb0284'
 
 def test_query_with_no_q_50(indDocSearch):
-    indDocSearch.query(collection='JUUL labs Collection', case='State of North Carolina', doc_type='email', n=50)
+    indDocSearch.query(industry='tobacco', collection='JUUL labs Collection', case='State of North Carolina', doc_type='email', n=50)
     assert len(indDocSearch.results) == 50
+    assert len(set([x['id'] for x in indDocSearch.results])) == 50
 
 def test_query_with_no_q_1000(indDocSearch):
-    indDocSearch.query(collection='JUUL labs Collection', case='State of North Carolina', doc_type='email', n=1000)
+    indDocSearch.query(industry='tobacco', collection='JUUL labs Collection', case='State of North Carolina', doc_type='email', n=1000)
     assert len(indDocSearch.results) == 1000
+    assert len(set([x['id'] for x in indDocSearch.results])) == 1000
 
 def test_save_parquet(indDocSearch, mock_results, tmp_path):
     indDocSearch.results = mock_results