Update get_rules_in_window.py

modules/get_rules_in_window.py

@@ -47,6 +47,8 @@ def get_date_range(start_date: str, end_mmdd: str = "01-03"):
 
 
 def get_rules(date_range: dict) -> list[dict]:
+    """Get rules within a date range.
+    """    
     results, _ = get_documents_by_date(
         start_date=date_range.get("start"), 
         end_date=date_range.get("end"), 
@@ -90,7 +92,16 @@ def filter_new_admin_rules(
         transition_year: int, 
         date_col: str = "publication_date", 
     ):
-    
+    """Remove rules issued by the new administration.
+
+    Args:
+        df (DataFrame): Input data.
+        transition_year (int): The year of the presidential transition.
+        date_col (str, optional): Column containing date information. Defaults to "publication_date".
+
+    Returns:
+        DataFrame: Filtered data.
+    """    
     admin_transitions = {
         2001: "george-w-bush", 
         2009: "barack-obama", 
@@ -139,7 +150,19 @@ def filter_corrections(df: DataFrame):
         raise SearchError(f"{len(df)} != {len(df_no_corrections)} + {len(df_corrections)}")
 
 
-def get_significant_rules(df, start_date):
+def get_significant_rules(df: DataFrame, start_date: str) -> tuple[DataFrame, date]:
+    """Get significant rules and merge with FR data.
+
+    Args:
+        df (DataFrame): Input data.
+        start_date (str): Start date of significant rule data.
+
+    Raises:
+        DataAvailabilityError: Raised when requesting significant rule counts prior to Executive Order 14094 of April 6, 2023.
+
+    Returns:
+        tuple[DataFrame, datetime.date]: Data with significant rules, last updated date for significant data
+    """
     process_columns = ("significant", "3f1_significant", )
     if date.fromisoformat(start_date) < date(2023, 4, 6):
         raise DataAvailabilityError("This program does not calculate significant rule counts prior to Executive Order 14094 of April 6, 2023.")
@@ -159,14 +182,24 @@ def get_significant_rules(df, start_date):
     return df, last_updated
 
 
-def get_rules_in_window(start_date: str, get_significant: bool = True, metadata=METADATA):
+def get_rules_in_window(start_date: str, get_significant: bool = True, metadata: dict = METADATA):
+    """Retrieve and process rules in a given CRA window.
+
+    Args:
+        start_date (str): Start date of window.
+        get_significant (bool, optional): Get significant rule data. Defaults to True.
+        metadata (dict, optional): Agency metadata. Defaults to METADATA.
+
+    Returns:
+        tuple[DataFrame, datetime.date]: Data with significant rules, last updated date for significant data
+    """    
     date_range = get_date_range(start_date)
     transition_year = date_range.get("transition_year")
     results = get_rules(date_range)
     df = format_documents(results)
     df, _ = filter_corrections(df)
     df = filter_new_admin_rules(df, transition_year)
-    df.loc[:, "acronym"] = get_agency_metadata_values(df, "parent_slug", metadata=METADATA, metadata_value="acronym")
+    df.loc[:, "acronym"] = get_agency_metadata_values(df, "parent_slug", metadata=metadata, metadata_value="acronym")
     if get_significant:
         df, last_updated = get_significant_rules(df, start_date)
     else:
@@ -174,12 +207,23 @@ def get_rules_in_window(start_date: str, get_significant: bool = True, metadata=
     return df, last_updated
 
 
-def get_list_agencies(start_date, agency_column: str = "parent_slug", significant: bool = True):
-    df, _ = get_rules_in_window(start_date, get_significant=significant)
+def get_list_agencies(start_date: str, agency_column: str = "parent_slug", significant: bool = True, **kwargs):
+    """Get list of agencies with rules in dataset.
+
+    Args:
+        start_date (str): Start date of window.
+        agency_column (str, optional): Column containing agency values. Defaults to "parent_slug".
+        significant (bool, optional): Get significant rule data. Defaults to True.
+
+    Returns:
+        list: List of agencies
+    """    
+    df, _ = get_rules_in_window(start_date, get_significant=significant, **kwargs)
     df_ex = df.explode(agency_column, ignore_index=True)
     return sorted(df_ex[agency_column].value_counts().index.to_list())
 
 
+# create objects to import in app
 DF, LAST_UPDATED = get_rules_in_window(START_DATE, get_significant=GET_SIGNIFICANT)
 AGENCIES = get_list_agencies(START_DATE, significant=GET_SIGNIFICANT)