readabs.grab_abs_url

  1"""Find and extract DataFrames from an ABS webpage."""
  2
  3# --- imports ---
  4# standard library imports
  5import zipfile
  6from functools import cache
  7from io import BytesIO
  8from typing import Any, Unpack
  9from pathlib import Path
 10
 11# analytic imports
 12import pandas as pd
 13from pandas import DataFrame
 14
 15from readabs.abs_catalogue import abs_catalogue
 16from readabs.download_cache import get_file
 17
 18# local imports
 19from readabs.get_abs_links import get_abs_links, get_table_name
 20from readabs.read_support import HYPHEN, ReadArgs, check_kwargs, get_args
 21
 22# --- constants ---
 23# File extensions for ABS data files
 24ZIP_EXTENSION = ".zip"
 25EXCEL_EXTENSION = ".xlsx"
 26
 27# Processing order: ZIP files must be processed before Excel files
 28# This prevents duplicate data when ZIP files contain Excel files
 29FILE_EXTENSIONS_PROCESSING_ORDER = (ZIP_EXTENSION, EXCEL_EXTENSION)
 30
 31# Default values and limits
 32EMPTY_BYTES_LENGTH = 0
 33
 34
 35# --- public - primary entry point for this module
 36@cache  # minimise slowness with repeat business
 37def grab_abs_url(
 38    cat: str = "",
 39    url: str = "",
 40    **kwargs: Unpack[ReadArgs],
 41) -> dict[str, DataFrame]:
 42    """For a given URL, extract the data from the Excel and ZIP file links found on that page.
 43
 44    The data is returned as a dictionary of DataFrames. The Excel files are converted
 45    into DataFrames, with each sheet in each Excel file becoming a separate DataFrame.
 46    ZIP files are examined for Excel files, which are similarly converted into
 47    DataFrames. The dictionary of DataFrames is returned.
 48
 49    The preferred mechanism for reading ABS data is to use the `read_abs_cat()`
 50    or `read_abs_series()` functions. This function is provided for those
 51    cases where the data is not available in the ABS catalogue, where the
 52    data is not a timeseries, or where the user wants to extract data from
 53    a specific ABS landingpage.
 54
 55
 56    Parameters
 57    ----------
 58    url : str = ""
 59        A URL for an ABS Catalogue landing page. Either a url or
 60        a catalogue number must be provided. If both are provided, the
 61        URL will be used.
 62
 63    cat : str = ""
 64        An ABS Catalogue number. If provided, and the URL is not
 65        provided, then the Catalogue number will be used to get the URL.
 66
 67    **kwargs : Unpack[ReadArgs]
 68        Accepts the same keyword arguments as `read_abs_cat()`.
 69
 70    Returns
 71    -------
 72    dict[str, DataFrame]
 73        A dictionary of DataFrames.
 74
 75    """
 76    # check/get the keyword arguments
 77    url = _get_url(url, cat)
 78    check_kwargs(kwargs, "grab_abs_url")  # warn if invalid kwargs
 79    args = get_args(kwargs, "grab_abs_url")  # get the valid kwargs
 80    if verbose := args["verbose"]:
 81        print(f"grab_abs_url(): {url=}, {args=}")
 82
 83    # get the URL links to the relevant ABS data files on that webpage
 84    links = get_abs_links(url, **args)
 85    if not links:
 86        print(f"No data files found at URL: {url}")
 87        return {}  # return an empty Dictionary
 88
 89    # read the data files into a dictionary of DataFrames
 90    abs_dict: dict[str, DataFrame] = {}
 91
 92    # Process single file requests first
 93    abs_dict = _process_single_files(abs_dict, links, args, verbose=verbose)
 94    if abs_dict:  # If single file was found and processed, return it
 95        return abs_dict
 96
 97    # Process all files based on configuration
 98    return _process_all_files(abs_dict, links, args)
 99
100
101def grab_abs_zip(
102    zip_path: Path | str,
103    **kwargs: Unpack[ReadArgs]
104) -> dict[str, DataFrame]:
105    """Grab and process a single ABS ZIP file from a file system location.
106
107    This is a convenience function that opens an ABS ZIP file from a local
108    filesystem path. Expect to be used rarely.
109
110    Parameters
111    ----------
112    zip_path : Path | str
113        The local filesystem path of the ABS ZIP file to open and process.
114
115    **kwargs : Unpack[ReadArgs]
116        Additional keyword arguments for file retrieval and processing.
117
118    Returns
119    -------
120    dict[str, DataFrame]
121        A dictionary of DataFrames extracted from the ZIP file.
122
123    """
124    check_kwargs(kwargs, "grab_abs_zip")  # warn if invalid kwargs
125    args = get_args(kwargs, "grab_abs_zip")  # get the valid kwargs
126
127    zp: Path = zip_path if isinstance(zip_path, Path) else Path(zip_path)
128    zip_bytes = zp.read_bytes()
129    abs_dict: dict[str, DataFrame] = {}
130    return _process_zip(abs_dict, zip_bytes, **args)
131
132
133# --- private
134def _process_single_files(
135    abs_dict: dict[str, DataFrame],
136    links: dict[str, list[str]],
137    args: dict[str, Any],  # ReadArgs after processing
138    *,
139    verbose: bool,
140) -> dict[str, DataFrame]:
141    """Process single file requests (single_excel_only or single_zip_only)."""
142    if args["single_excel_only"]:
143        link = _find_url(links, EXCEL_EXTENSION, args["single_excel_only"], verbose=verbose)
144        if link:
145            return _add_excel(abs_dict, link, **args)
146
147    if args["single_zip_only"]:
148        link = _find_url(links, ZIP_EXTENSION, args["single_zip_only"], verbose=verbose)
149        if link:
150            return _add_zip(abs_dict, link, **args)
151
152    return abs_dict
153
154
155def _process_all_files(
156    abs_dict: dict[str, DataFrame],
157    links: dict[str, list[str]],
158    args: dict[str, Any],  # ReadArgs after processing
159) -> dict[str, DataFrame]:
160    """Process all files based on configuration (get_zip, get_excel, etc.)."""
161    for link_type in FILE_EXTENSIONS_PROCESSING_ORDER:
162        for link in links.get(link_type, []):
163            if link_type == ZIP_EXTENSION and args["get_zip"]:
164                abs_dict = _add_zip(abs_dict, link, **args)
165            elif link_type == EXCEL_EXTENSION and _should_process_excel_file(args, links):
166                abs_dict = _add_excel(abs_dict, link, **args)
167    return abs_dict
168
169
170def _should_process_excel_file(args: dict[str, Any], links: dict[str, list[str]]) -> bool:
171    """Determine if Excel files should be processed based on configuration.
172
173    Excel files are processed if:
174    1. get_excel is explicitly True, or
175    2. get_excel_if_no_zip is True AND (get_zip is False OR no ZIP files are available)
176
177    Args:
178        args: Configuration arguments from user
179        links: Dictionary of available file links by type
180
181    Returns:
182        bool: True if Excel files should be processed
183
184    """
185    # Always process if explicitly requested
186    if args["get_excel"]:
187        return True
188
189    # Process Excel if requested when no ZIP files, and either:
190    # - ZIP processing is disabled, or
191    # - No ZIP files are available
192    if args["get_excel_if_no_zip"]:
193        zip_processing_disabled = not args["get_zip"]
194        no_zip_files_available = not links.get(ZIP_EXTENSION, [])
195        return zip_processing_disabled or no_zip_files_available
196
197    return False
198
199
200def _find_url(links: dict[str, list[str]], targ_type: str, target: str, *, verbose: bool = False) -> str:
201    """Find the URL for a target file type.
202
203    Args:
204        links: Dictionary mapping file types to lists of URLs
205        targ_type: Target file extension (e.g., '.xlsx', '.zip')
206        target: Target filename without extension
207        verbose: Whether to print debug information
208
209    Returns:
210        str: The matching URL if found, otherwise an empty string
211
212    """
213    targ_list = links.get(targ_type, [])
214    if not targ_list:
215        return ""
216    goal = f"{target}{targ_type}"
217    if verbose:
218        print(f"_find_url(): looking for {goal} in {targ_list}.")
219    for link in targ_list:
220        if link.endswith(goal):
221            return link
222    return ""
223
224
225def _get_url(url: str, cat: str) -> str:
226    """Get URL from provided URL or catalogue number.
227
228    If an ABS catalogue number is provided and URL is not provided,
229    get the URL for the ABS data files on the ABS webpage.
230    Otherwise, return the URL provided. Either the 'url' or
231    'cat' argument must be provided.
232
233    Args:
234        url: The URL to use if provided
235        cat: The catalogue number to use if URL is not provided
236
237    Returns:
238        str: The URL to use for data retrieval
239
240    Raises:
241        ValueError: If neither URL nor valid catalogue number is provided
242
243    """
244    if not url and cat:
245        try:
246            cat_map = abs_catalogue()
247            if cat in cat_map.index:
248                url = str(cat_map.loc[cat, "URL"])
249        except (KeyError, IndexError) as e:
250            raise ValueError(f"Catalogue number '{cat}' not found in ABS catalogue: {e}") from e
251        except (ConnectionError, TimeoutError) as e:
252            raise ValueError(f"Network error retrieving catalogue for '{cat}': {e}") from e
253        except (ValueError, TypeError) as e:
254            raise ValueError(f"Invalid catalogue data for '{cat}': {e}") from e
255
256    if not url:
257        raise ValueError("_get_url(): no URL or valid catalogue number provided.")
258
259    return url
260
261
262def _process_zip(
263    abs_dict: dict[str, DataFrame],
264    zip_contents: bytes,
265    **args: Any,  # ReadArgs compatible
266) -> dict[str, DataFrame]:
267    """Read and process a ZIP file's contents from bytes."""
268
269    if len(zip_contents) == EMPTY_BYTES_LENGTH:
270        return abs_dict
271
272    with zipfile.ZipFile(BytesIO(zip_contents)) as zipped:
273        for element in zipped.infolist():
274            # get the zipfile into pandas
275            table_name = get_table_name(url=element.filename)
276            raw_bytes = zipped.read(element.filename)
277            abs_dict = _add_excel_bytes(abs_dict, raw_bytes, table_name, args)
278
279    return abs_dict
280
281
282def _add_zip(
283    abs_dict: dict[str, DataFrame],
284    link: str,
285    **args: Any,  # ReadArgs compatible
286) -> dict[str, DataFrame]:
287    """Read and process a ZIP file from a URL.
288
289    Downloads the ZIP file and iterates over its contents, calling
290    _add_excel_bytes() to extract Excel files and add their contents
291    to the DataFrames dictionary.
292
293    Args:
294        abs_dict: Dictionary to store extracted DataFrames
295        link: URL to the ZIP file
296        **args: Additional arguments passed to file retrieval functions
297
298    Returns:
299        dict[str, DataFrame]: Updated dictionary with new DataFrames from ZIP contents
300
301    """
302    zip_contents = get_file(link, **args)
303    return _process_zip(abs_dict, zip_contents, **args)
304
305
306def _add_excel_bytes(
307    abs_dict: dict[str, DataFrame],
308    raw_bytes: bytes,
309    name: str,
310    args: dict[str, Any],  # ReadArgs after processing
311) -> dict[str, DataFrame]:
312    """Convert Excel file bytes to DataFrames and add to dictionary.
313
314    Processes the bytes as an Excel file, converting each sheet to a DataFrame
315    and adding them to the dictionary using 'name---sheet_name' as keys.
316
317    Args:
318        abs_dict: Dictionary to store extracted DataFrames
319        raw_bytes: Bytes content of the Excel file
320        name: Base name for the Excel file
321        args: Dictionary of processing arguments
322
323    Returns:
324        dict[str, DataFrame]: Updated dictionary with new DataFrames from Excel sheets
325
326    """
327    verbose = args.get("verbose", False)
328
329    if len(raw_bytes) == EMPTY_BYTES_LENGTH:
330        if verbose:
331            print("_add_excel_bytes(): the raw bytes are empty.")
332        return abs_dict
333
334    # convert the raw bytes into a pandas ExcelFile
335    try:
336        excel = pd.ExcelFile(BytesIO(raw_bytes))
337    except (ValueError, TypeError) as e:
338        message = f"With {name}: could not convert raw bytes to ExcelFile.\n{e}"
339        print(message)
340        return abs_dict
341
342    # iterate over the sheets in the Excel file
343    for sheet_name in excel.sheet_names:
344        # grab and go - no treatment of the data
345        sheet_data = excel.parse(
346            sheet_name,
347        )
348        if len(sheet_data) == EMPTY_BYTES_LENGTH:
349            if verbose:
350                print(f"_add_excel_bytes(): sheet {sheet_name} in {name} is empty.")
351            continue
352        abs_dict[f"{name}{HYPHEN}{sheet_name}"] = sheet_data
353
354    # return the dictionary of DataFrames
355    return abs_dict
356
357
358def _add_excel(
359    abs_dict: dict[str, DataFrame],
360    link: str,
361    **args: Any,  # ReadArgs compatible
362) -> dict[str, DataFrame]:
363    """Read in an Excel file at the URL in the 'link' argument.
364
365    Pass those bytes to _add_excel_bytes() to put the contents
366    into the dictionary of DataFrames given by 'abs_dict'. When done,
367    return the dictionary of DataFrames.
368    """
369    name = get_table_name(link)
370
371    if name in abs_dict:
372        # table already in the dictionary
373        return abs_dict
374
375    raw_bytes = get_file(link, **args)
376
377    return _add_excel_bytes(abs_dict, raw_bytes, name, args)
378
379
380# --- main ---
381if __name__ == "__main__":
382
383    def simple_test() -> None:
384        """Test the grab_abs_url and grab_abs_zip functions."""
385
386        def test(name: str, **kwargs: Any) -> None:  # ReadArgs compatible
387            print(f"TEST -- {name}")
388            try:
389                data_dict = grab_abs_url(**kwargs)
390                print("---")
391                if not data_dict:
392                    print("PROBLEM -- No data found.")
393                else:
394                    print(f"SUCCESS -- Found {len(data_dict)} datasets: {list(data_dict.keys())}")
395            except Exception as e:  # pylint: disable=broad-except
396                print(f"ERROR -- Test failed with exception: {e}")
397            print(f"Done.\n{'=' * 20}\n")
398
399        def test_zip(zip_path: Path, **kwargs: Any) -> None:  # ReadArgs compatible
400            print(f"TEST -- grab_abs_zip() with {zip_path}")
401            try:
402                data_dict = grab_abs_zip(zip_path, **kwargs)
403                print("---")
404                if not data_dict:
405                    print("PROBLEM -- No data found.")
406                else:
407                    print(f"SUCCESS -- Found {len(data_dict)} datasets: {list(data_dict.keys())}")
408            except Exception as e:  # pylint: disable=broad-except
409                print(f"ERROR -- Test failed with exception: {e}")
410            print(f"Done.\n{'=' * 20}\n")
411
412        # 4 -- grab from url
413        _name = ".test-data/Qrtly-CPI-Time-series-spreadsheets-all.zip"
414        _zip_path = Path(_name)
415        test_zip(_zip_path, verbose=True)
416
417        # --- various grab_abs_url() tests
418
419        name = "1 -- grab a single zip file"
420        test(
421            name,
422            cat="6291.0.55.001",
423            single_zip_only="p6291_all_quarterly_spreadsheets",
424            get_zip=True,
425            verbose=True,
426        )
427
428        name = "2 -- grab a single Excel file"
429        test(
430            name,
431            cat="6202.0",
432            get_excel=True,
433            single_excel_only="6202001",
434            verbose=False,
435        )
436
437        # 3 -- grab the whole shebang
438        urls = [
439            "https://www.abs.gov.au/statistics/labour/jobs/weekly-payroll-jobs/latest-release",
440            "https://www.abs.gov.au/statistics/people/population/national-state-and-territory-population/dec-2023",
441        ]
442        for i, url_ in enumerate(urls):
443            name = f"3.{i} -- grab the whole shebang {url_}"
444            test(name, url=url_, verbose=True)
445
446    simple_test()