Lint and format

si2-urssi · Feb 26, 2024 · cc95f54 · cc95f54
1 parent 94ba108
commit cc95f54
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 14 deletions.
diff --git a/award_pynder/sources/mellon.py b/award_pynder/sources/mellon.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python
 
 from __future__ import annotations
-from copy import deepcopy
 
 import logging
+from copy import deepcopy
 from datetime import datetime
 
 import pandas as pd
@@ -52,7 +52,7 @@
     totalCount
 }
 """.strip()  # noqa: E501
-_DEFAULT_BULK_PARAMS = {
+_DEFAULT_BULK_PARAMS: dict = {
     "operationName": "GrantFilterQuery",
     "variables": {
         "limit": _DEFAULT_CHUNK_SIZE,
@@ -82,7 +82,7 @@
     }
 }
 """.strip()
-_DEFAULT_SINGLE_PARAMS = {
+_DEFAULT_SINGLE_PARAMS: dict = {
     "variables": {
         "grantId": "",
     },
@@ -119,7 +119,7 @@ def _format_query_params(
                     to_datetime_parsed.year + 1,
                 )
             )
-        
+
         # Copy and update params
         params = deepcopy(_DEFAULT_BULK_PARAMS)
         params["variables"]["term"] = query or ""
@@ -129,10 +129,10 @@ def _format_query_params(
 
         # Return the full API string
         return params
-    
+
     @staticmethod
     def _query_total_grants(
-        query: str| None = None,
+        query: str | None = None,
         from_datetime: str | datetime | None = None,
         to_datetime: str | datetime | None = None,
     ) -> int:
@@ -194,7 +194,7 @@ def _format_dataframe(
 
         # Create new dataframe with only the columns we want
         return df[ALL_DATASET_FIELDS]
-    
+
     @staticmethod
     def _get_funded_amount_for_grant(grant_id: str) -> float:
         # Copy and update params
@@ -214,7 +214,7 @@ def _get_funded_amount_for_grant(grant_id: str) -> float:
 
             # Process results
             return data["data"]["grantDetails"]["grant"]["amount"]
-        
+
         except Exception as e:
             raise ValueError(f"Error while fetching grant amount: {e}") from e
 
@@ -247,9 +247,9 @@ def _get_chunk(
             data = resp.json()
 
             # Process results
-            chunk = pd.DataFrame([
-                item["data"] for item in data["data"]["grantSearch"]["entities"]
-            ])
+            chunk = pd.DataFrame(
+                [item["data"] for item in data["data"]["grantSearch"]["entities"]]
+            )
 
             # Make API calls for amount funded
             chunk["amount"] = chunk["id"].apply(Mellon._get_funded_amount_for_grant)
@@ -316,5 +316,5 @@ def get_data(
         #         (all_results["dateAwarded"] >= from_datetime)
         #         & (all_results["dateAwarded"] <= to_datetime)
         #     ]
-        
-        return all_results
+
+        return all_results
diff --git a/award_pynder/tests/sources/test_mellon.py b/award_pynder/tests/sources/test_mellon.py
@@ -9,7 +9,12 @@
 
 def test_nsf() -> None:
     # Get data
-    df = Mellon.get_data(query="software", from_datetime="2023-01-01", to_datetime="2023-06-01", tqdm_kwargs={"leave": False})
+    df = Mellon.get_data(
+        query="software",
+        from_datetime="2023-01-01",
+        to_datetime="2023-06-01",
+        tqdm_kwargs={"leave": False},
+    )
 
     # Run tests
     assert_dataset_basics(df)