Finesse new URL Probe logic

src/core/env_var_manager.py

@@ -16,7 +16,8 @@ def __init__(self, env: dict = os.environ):
         self.env = env
         self._load()
 
-    def _load(self):
+    def _load(self) -> None:
+        """Load environment variables from environment"""
 
         self.google_api_key = self.require_env("GOOGLE_API_KEY")
         self.google_cse_id = self.require_env("GOOGLE_CSE_ID")

src/core/tasks/scheduled/loader.py

@@ -41,6 +41,14 @@ async def load_entries(self) -> list[ScheduledTaskEntry]:
 
 
         return [
+            ScheduledTaskEntry(
+                operator=SyncDataSourcesTaskOperator(
+                    adb_client=self.async_core.adb_client,
+                    pdap_client=self.pdap_client
+                ),
+                interval=IntervalEnum.DAILY,
+                enabled=self.env.bool("SYNC_DATA_SOURCES_TASK_FLAG", default=True)
+            ),
             ScheduledTaskEntry(
                 operator=RunURLTasksTaskOperator(async_core=self.async_core),
                 interval=IntervalEnum.HOURLY,
@@ -57,14 +65,6 @@ async def load_entries(self) -> list[ScheduledTaskEntry]:
                 interval=IntervalEnum.DAILY,
                 enabled=self.env.bool("POPULATE_BACKLOG_SNAPSHOT_TASK_FLAG", default=True)
             ),
-            ScheduledTaskEntry(
-                operator=SyncDataSourcesTaskOperator(
-                    adb_client=self.async_core.adb_client,
-                    pdap_client=self.pdap_client
-                ),
-                interval=IntervalEnum.DAILY,
-                enabled=self.env.bool("SYNC_DATA_SOURCES_TASK_FLAG", default=True)
-            ),
             ScheduledTaskEntry(
                 operator=SyncAgenciesTaskOperator(
                     adb_client=self.async_core.adb_client,

src/core/tasks/url/operators/html/scraper/parser/core.py

@@ -35,7 +35,12 @@
             html_info: ResponseHTMLInfo,
             parser_type: ParserTypeEnum,
             html_content: str
-    ):
+    ) -> None:
+        """
+        Modifies:
+            html_info
+        """
+
         soup = BeautifulSoup(
             markup=html_content,
             features=parser_type.value,
@@ -48,7 +53,7 @@
         if soup.html is not None:
             soup.html.decompose()
 
-    def get_div_text(self, soup):
+    def get_div_text(self, soup: BeautifulSoup) -> str:
         div_text = ""
         MAX_WORDS = 500
         for div in soup.find_all("div"):
@@ -85,29 +90,43 @@
                 continue
             setattr(html_info, header_tag, tag_content)
 
-    def get_html_title(self, soup: BeautifulSoup) -> Optional[str]:
+    def get_html_title(self, soup: BeautifulSoup) -> str | None:
         if soup.title is None:
             return None
         if soup.title.string is None:
             return None
         return remove_excess_whitespace(soup.title.string)
 
 
-    def add_url_and_path(self, html_info: ResponseHTMLInfo, html_content: str, url: str):
+    def add_url_and_path(
+        self,
+        html_info: ResponseHTMLInfo,
+        html_content: str,
+        url: str
+    ) -> None:
+        """
+        Modifies:
+            html_info.url
+            html_info.url_path
+        """
         url = add_https(url)
         html_info.url = url
 
         url_path = drop_hostname(url)
         url_path = remove_trailing_backslash(url_path)
         html_info.url_path = url_path
 
-    async def add_root_page_titles(self, html_info: ResponseHTMLInfo):
+    async def add_root_page_titles(self, html_info: ResponseHTMLInfo) -> None:
+        """
+        Modifies:
+            html_info.root_page_title
+        """
         root_page_title = await self.root_url_cache.get_title(html_info.url)
         html_info.root_page_title = remove_excess_whitespace(
             root_page_title
         )
 
-    def get_parser_type(self, content_type: str) -> ParserTypeEnum or None:
+    def get_parser_type(self, content_type: str) -> ParserTypeEnum | None:
         try:
             # If content type does not contain "html" or "xml" then we can assume that the content is unreadable
             if "html" in content_type:

src/core/tasks/url/operators/html/scraper/parser/util.py

@@ -5,7 +5,9 @@
 from src.core.tasks.url.operators.html.scraper.parser.dtos.response_html import ResponseHTMLInfo
 
 
-def convert_to_response_html_info(html_content_infos: list[URLHTMLContentInfo]):
+def convert_to_response_html_info(
+    html_content_infos: list[URLHTMLContentInfo]
+) -> ResponseHTMLInfo:
     response_html_info = ResponseHTMLInfo()
 
     for html_content_info in html_content_infos:
@@ -32,12 +34,12 @@
     return url
 
 
-def remove_trailing_backslash(url_path):
+def remove_trailing_backslash(url_path: str) -> str:
     if url_path and url_path[-1] == "/":
         url_path = url_path[:-1]
     return url_path
 
 
-def drop_hostname(new_url):
+def drop_hostname(new_url: str) -> str:
     url_path = urlparse(new_url).path[1:]
     return url_path

src/core/tasks/url/operators/html/scraper/root_url_cache/core.py

@@ -12,19 +12,19 @@
 
 
 class RootURLCache:
-    def __init__(self, adb_client: Optional[AsyncDatabaseClient] = None):
+    def __init__(self, adb_client: AsyncDatabaseClient | None = None):
         if adb_client is None:
             adb_client = AsyncDatabaseClient()
         self.adb_client = adb_client
         self.cache = None
 
-    async def save_to_cache(self, url: str, title: str):
+    async def save_to_cache(self, url: str, title: str) -> None:
         if url in self.cache:
             return
         self.cache[url] = title
         await self.adb_client.add_to_root_url_cache(url=url, page_title=title)
 
-    async def get_from_cache(self, url: str) -> Optional[str]:
+    async def get_from_cache(self, url: str) -> str | None:
         if self.cache is None:
             self.cache = await self.adb_client.load_root_url_cache()
 

src/core/tasks/url/operators/misc_metadata/core.py

@@ -22,16 +22,16 @@
         super().__init__(adb_client)
 
     @property
-    def task_type(self):
+    def task_type(self) -> TaskType:
         return TaskType.MISC_METADATA
 
-    async def meets_task_prerequisites(self):
+    async def meets_task_prerequisites(self) -> bool:
         return await self.adb_client.has_pending_urls_missing_miscellaneous_metadata()
 
     async def get_subtask(
             self,
             collector_type: CollectorType
-    ) -> Optional[MiscellaneousMetadataSubtaskBase]:
+    ) -> MiscellaneousMetadataSubtaskBase | None:
         match collector_type:
             case CollectorType.MUCKROCK_SIMPLE_SEARCH:
                 return MuckrockMiscMetadataSubtask()
@@ -47,12 +47,17 @@
                 return None
 
     async def html_default_logic(self, tdo: URLMiscellaneousMetadataTDO):
+        """
+        Modifies:
+            tdo.name
+            tdo.description
+        """
         if tdo.name is None:
             tdo.name = tdo.html_metadata_info.title
         if tdo.description is None:
             tdo.description = tdo.html_metadata_info.description
 
-    async def inner_task_logic(self):
+    async def inner_task_logic(self) -> None:
         tdos: list[URLMiscellaneousMetadataTDO] = await self.adb_client.get_pending_urls_missing_miscellaneous_metadata()
         await self.link_urls_to_task(url_ids=[tdo.url_id for tdo in tdos])
 

src/core/tasks/url/operators/probe_404/core.py

@@ -26,13 +26,17 @@
         self.url_request_interface = url_request_interface
 
     @property
-    def task_type(self):
+    def task_type(self) -> TaskType:
         return TaskType.PROBE_404
 
-    async def meets_task_prerequisites(self):
+    async def meets_task_prerequisites(self) -> bool:
         return await self.adb_client.has_pending_urls_not_recently_probed_for_404()
 
-    async def probe_urls_for_404(self, tdos: list[URL404ProbeTDO]):
+    async def probe_urls_for_404(self, tdos: list[URL404ProbeTDO]) -> None:
+        """
+        Modifies:
+            URL404ProbeTDO.is_404
+        """
         responses = await self.url_request_interface.make_simple_requests(
             urls=[tdo.url for tdo in tdos]
         )
@@ -42,7 +46,7 @@
             tdo.is_404 = response.status == HTTPStatus.NOT_FOUND
 
 
-    async def inner_task_logic(self):
+    async def inner_task_logic(self) -> None:
         tdos = await self.get_pending_urls_not_recently_probed_for_404()
         url_ids = [task_info.url_id for task_info in tdos]
         await self.link_urls_to_task(url_ids=url_ids)
@@ -55,9 +59,17 @@
     async def get_pending_urls_not_recently_probed_for_404(self) -> list[URL404ProbeTDO]:
         return await self.adb_client.get_pending_urls_not_recently_probed_for_404()
 
-    async def update_404s_in_database(self, url_ids_404: list[int]):
+    async def update_404s_in_database(self, url_ids_404: list[int]) -> None:
+        """
+        Modifies:
+            URL data in DB
+        """
         await self.adb_client.mark_all_as_404(url_ids_404)
 
-    async def mark_as_recently_probed_for_404(self, url_ids: list[int]):
+    async def mark_as_recently_probed_for_404(self, url_ids: list[int]) -> None:
+        """
+        Modifies:
+            URL data in DB
+        """
         await self.adb_client.mark_all_as_recently_probed_for_404(url_ids)
 

src/db/helpers/connect.py

@@ -1,5 +1,5 @@
 from src.core.env_var_manager import EnvVarManager
 
 
-def get_postgres_connection_string(is_async = False):
+def get_postgres_connection_string(is_async = False) -> str:
     return EnvVarManager.get().get_postgres_connection_string(is_async)

src/db/models/helpers.py

@@ -1,7 +1,7 @@
 from sqlalchemy import Column, TIMESTAMP, func, Integer, ForeignKey, Enum as SAEnum
 from enum import Enum as PyEnum
 
-def get_created_at_column():
+def get_created_at_column() -> Column:
     return Column(TIMESTAMP, nullable=False, server_default=CURRENT_TIME_SERVER_DEFAULT)
 
 

src/db/queries/implementations/core/common/annotation_exists.py

@@ -41,7 +41,7 @@ def get_all(self) -> list[Any]:
 
     async def _annotation_exists_case(
         self,
-    ):
+    ) -> list[Any]:
         cases = []
         for model in ALL_ANNOTATION_MODELS:
             cases.append(

src/external/huggingface/hub/client.py

@@ -11,10 +11,21 @@ class HuggingFaceHubClient:
     def __init__(self, token: str):
         self.token = token
 
-    def _push_dataset_to_hub(self, repo_id: str, dataset: Dataset):
+    def _push_dataset_to_hub(self, repo_id: str, dataset: Dataset) -> None:
+        """
+        Modifies:
+            - repository on Hugging Face, identified by `repo_id`
+        """
         dataset.push_to_hub(repo_id=repo_id, token=self.token)
 
-    def push_data_sources_raw_to_hub(self, outputs: list[GetForLoadingToHuggingFaceOutput]):
+    def push_data_sources_raw_to_hub(
+        self,
+        outputs: list[GetForLoadingToHuggingFaceOutput]
+    ) -> None:
+        """
+        Modifies:
+            - repository on Hugging Face, identified by `DATA_SOURCES_RAW_REPO_ID`
+        """
         dataset = format_as_huggingface_dataset(outputs)
         print(dataset)
         self._push_dataset_to_hub(repo_id=DATA_SOURCES_RAW_REPO_ID, dataset=dataset)

src/external/pdap/client.py

@@ -192,14 +192,15 @@ async def sync_data_sources(
         )
         headers = await self.access_manager.jwt_header()
         headers['Content-Type'] = "application/json"
+        params_dict = {"page": params.page}
+        if params.cutoff_date is not None:
+            params_dict["updated_at"] = params.cutoff_date
+
         request_info = RequestInfo(
             type_=RequestType.GET,
             url=url,
             headers=headers,
-            params={
-                "page": params.page,
-                "updated_at": params.cutoff_date
-            }
+            params=params_dict
         )
         response_info = await self.access_manager.make_request(request_info)
         return DataSourcesSyncResponseInfo(

src/external/pdap/dtos/match_agency/post.py

@@ -6,6 +6,6 @@
 class MatchAgencyInfo(BaseModel):
     id: int
     submitted_name: str
-    state: Optional[str] = None
-    county: Optional[str] = None
-    locality: Optional[str] = None
+    state: str | None = None
+    county: str | None = None
+    locality: str | None = None

src/external/pdap/dtos/sync/agencies.py

@@ -6,9 +6,9 @@
 class AgenciesSyncResponseInnerInfo(BaseModel):
     display_name: str
     agency_id: int
-    state_name: Optional[str]
-    county_name: Optional[str]
-    locality_name: Optional[str]
+    state_name: str | None
+    county_name: str | None
+    locality_name: str | None
     updated_at: datetime.datetime
 
 class AgenciesSyncResponseInfo(BaseModel):

src/external/pdap/dtos/unique_url_duplicate.py

@@ -8,4 +8,4 @@
 class UniqueURLDuplicateInfo(BaseModel):
     original_url: str
     approval_status: ApprovalStatus
-    rejection_note: Optional[str] = None
+    rejection_note: str | None = None

src/external/url_request/core.py

@@ -2,7 +2,6 @@
 
 from src.external.url_request.dtos.url_response import URLResponseInfo
 from src.external.url_request.probe.core import URLProbeManager
-from src.external.url_request.probe.models.response import URLProbeResponse
 from src.external.url_request.probe.models.wrapper import URLProbeResponseOuterWrapper
 from src.external.url_request.request import fetch_urls
 

src/external/url_request/probe/convert.py

@@ -53,6 +53,7 @@ def _extract_destination_url(cr: ClientResponse) -> str:
     return str(cr.url)
 
 def convert_client_response_to_probe_response(
+    url: str,
     cr: ClientResponse
 ) -> URLProbeResponse | URLProbeRedirectResponsePair:
     error = _extract_error(cr)
@@ -69,13 +70,12 @@ def convert_client_response_to_probe_response(
     source_cr = cr.history[0]  # Source CR is the first in the history
     destination_cr = cr
 
-    source_url = str(source_cr.url)
     destination_url = str(destination_cr.url)
 
     source_error = _extract_error(source_cr)
     source_content_type = _extract_content_type(source_cr, error=source_error)
     source_probe_response = URLProbeResponse(
-        url=source_url,
+        url=url,
         status_code=source_cr.status,
         content_type=source_content_type,
         error=source_error,