Create utils.py

utils.py

@@ -0,0 +1,433 @@
+"""
+general utility functions for loading, saving, and manipulating data
+"""
+
+import os
+import logging
+import pprint as pp
+import re
+import shutil  # zipfile formats
+import warnings
+from datetime import datetime
+from os.path import basename, getsize, join
+from pathlib import Path
+import logging
+
+import pandas as pd
+import requests
+from natsort import natsorted
+from symspellpy import SymSpell
+from tqdm.auto import tqdm
+
+import warnings
+
+warnings.filterwarnings(
+    action="ignore", message=".*the GPL-licensed package `unidecode` is not installed*"
+)  # cleantext GPL-licensed package reminder is annoying
+
+
+class DisableLogger:
+    def __enter__(self):
+        logging.disable(logging.CRITICAL)
+
+    def __exit__(self, exit_type, exit_value, exit_traceback):
+        logging.disable(logging.NOTSET)
+
+
+with DisableLogger():
+    from cleantext import clean
+
+
+def clear_loggers():
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+
+
+def get_timestamp():
+    return datetime.now().strftime("%b-%d-%Y_t-%H")
+
+
+def print_spacer(n=1):
+    """print_spacer - print a spacer line"""
+    print("\n   --------    " * n)
+
+
+def remove_trailing_punctuation(text: str):
+    """
+    remove_trailing_punctuation - remove trailing punctuation from a string
+
+    Args:
+        text (str): [string to be cleaned]
+
+    Returns:
+        [str]: [cleaned string]
+    """
+    return text.strip("?!.,;:")
+
+
+def correct_phrase_load(my_string: str):
+    """
+    correct_phrase_load [basic / unoptimized implementation of SymSpell to correct a string]
+
+    Args:
+        my_string (str): [text to be corrected]
+
+    Returns:
+        str: the corrected string
+    """
+    sym_spell = SymSpell(max_dictionary_edit_distance=2, prefix_length=7)
+
+    dictionary_path = (
+        r"symspell_rsc/frequency_dictionary_en_82_765.txt"  # from repo root
+    )
+    bigram_path = (
+        r"symspell_rsc/frequency_bigramdictionary_en_243_342.txt"  # from repo root
+    )
+    # term_index is the column of the term and count_index is the
+    # column of the term frequency
+    sym_spell.load_dictionary(dictionary_path, term_index=0, count_index=1)
+    sym_spell.load_bigram_dictionary(bigram_path, term_index=0, count_index=2)
+
+    # max edit distance per lookup (per single word, not per whole input string)
+    suggestions = sym_spell.lookup_compound(
+        clean(my_string), max_edit_distance=2, ignore_non_words=True
+    )
+    if len(suggestions) < 1:
+        return my_string
+    else:
+        first_result = suggestions[0]
+        return first_result._term
+
+
+def fast_scandir(dirname: str):
+    """
+    fast_scandir [an os.path-based means to return all subfolders in a given filepath]
+
+    """
+
+    subfolders = [f.path for f in os.scandir(dirname) if f.is_dir()]
+    for dirname in list(subfolders):
+        subfolders.extend(fast_scandir(dirname))
+    return subfolders  # list
+
+
+def create_folder(directory: str):
+
+    os.makedirs(directory, exist_ok=True)
+
+
+def chunks(lst: list, n: int):
+    """
+    chunks   -  Yield successive n-sized chunks from lst
+    Args:   lst (list): list to be chunked
+    n (int): size of chunks
+
+    """
+
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]
+
+
+def shorten_list(
+    list_of_strings: list, max_chars: int = 512, no_blanks=True, verbose=False
+):
+    """a helper function that iterates through a list backwards, adding to a new list.
+
+        When <max_chars> is met, that list entry is not added.
+    Args:
+        list_of_strings (list): list of strings to be shortened
+        max_chars (int, optional): maximum number of characters in a the list in total. Defaults to 512.
+        no_blanks (bool, optional): if True, blank strings are not added to the new list. Defaults to True.
+        verbose (bool, optional): if True, print the list of strings before and after the shorten. Defaults to False.
+    """
+    list_of_strings = [
+        str(x) for x in list_of_strings
+    ]  # convert to strings if not already
+    shortened_list = []
+    total_len = 0
+    for i, string in enumerate(list_of_strings[::-1], start=1):
+
+        if len(string.strip()) == 0 and no_blanks:
+            continue
+        if len(string) + total_len >= max_chars:
+            logging.info(f"string # {i} puts total over limit, breaking ")
+            break
+        total_len += len(string)
+        shortened_list.insert(0, string)
+    if len(shortened_list) == 0:
+        logging.info(f"shortened list with max_chars={max_chars} has no entries")
+    if verbose:
+        print(f"total length of list is {total_len} chars")
+    return shortened_list
+
+
+def chunky_pandas(my_df, num_chunks: int = 4):
+    """
+    chunky_pandas [split dataframe into `num_chunks` equal chunks, return each inside a list]
+
+    Args:
+        my_df (pd.DataFrame)
+        num_chunks (int, optional): Defaults to 4.
+
+    Returns:
+        list: a list of dataframes
+    """
+    n = int(len(my_df) // num_chunks)
+    list_df = [my_df[i : i + n] for i in range(0, my_df.shape[0], n)]
+
+    return list_df
+
+
+def load_dir_files(
+    directory: str, req_extension=".txt", return_type="list", verbose=False
+):
+    """
+    load_dir_files - an os.path based method of returning all files with extension `req_extension` in a given directory and subdirectories
+
+    Args:
+
+
+    Returns:
+        list or dict: an iterable of filepaths or a dict of filepaths and their respective filenames
+    """
+    appr_files = []
+    # r=root, d=directories, f = files
+    for r, d, f in os.walk(directory):
+        for prefile in f:
+            if prefile.endswith(req_extension):
+                fullpath = os.path.join(r, prefile)
+                appr_files.append(fullpath)
+
+    appr_files = natsorted(appr_files)
+
+    if verbose:
+        print("A list of files in the {} directory are: \n".format(directory))
+        if len(appr_files) < 10:
+            pp.pprint(appr_files)
+        else:
+            pp.pprint(appr_files[:10])
+            print("\n and more. There are a total of {} files".format(len(appr_files)))
+
+    if return_type.lower() == "list":
+        return appr_files
+    else:
+        if verbose:
+            print("returning dictionary")
+
+        appr_file_dict = {}
+        for this_file in appr_files:
+            appr_file_dict[basename(this_file)] = this_file
+
+        return appr_file_dict
+
+
+def URL_string_filter(text):
+    """
+    URL_string_filter - filter out nonstandard "text" characters
+
+    """
+    custom_printable = (
+        "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ._"
+    )
+
+    filtered = "".join((filter(lambda i: i in custom_printable, text)))
+
+    return filtered
+
+
+def getFilename_fromCd(cd):
+    """getFilename_fromCd - get the filename from a given cd str"""
+    if not cd:
+        return None
+    fname = re.findall("filename=(.+)", cd)
+    if len(fname) > 0:
+        output = fname[0]
+    elif cd.find("/"):
+        possible_fname = cd.rsplit("/", 1)[1]
+        output = URL_string_filter(possible_fname)
+    else:
+        output = None
+    return output
+
+
+def get_zip_URL(
+    URLtoget: str,
+    extract_loc: str = None,
+    file_header: str = "dropboxexport_",
+    verbose: bool = False,
+):
+    """get_zip_URL - download a zip file from a given URL and extract it to a given location"""
+
+    r = requests.get(URLtoget, allow_redirects=True)
+    names = getFilename_fromCd(r.headers.get("content-disposition"))
+    fixed_fnames = names.split(";")  # split the multiple results
+    this_filename = file_header + URL_string_filter(fixed_fnames[0])
+
+    # define paths and save the zip file
+    if extract_loc is None:
+        extract_loc = "dropbox_dl"
+    dl_place = join(os.getcwd(), extract_loc)
+    create_folder(dl_place)
+    save_loc = join(os.getcwd(), this_filename)
+    open(save_loc, "wb").write(r.content)
+    if verbose:
+        print("downloaded file size was {} MB".format(getsize(save_loc) / 1000000))
+
+    # unpack the archive
+    shutil.unpack_archive(save_loc, extract_dir=dl_place)
+    if verbose:
+        print("extracted zip file - ", datetime.now())
+        x = load_dir_files(dl_place, req_extension="", verbose=verbose)
+
+    # remove original
+    try:
+        os.remove(save_loc)
+        del save_loc
+    except Exception:
+        print("unable to delete original zipfile - check if exists", datetime.now())
+
+    print("finished extracting zip - ", datetime.now())
+
+    return dl_place
+
+
+def merge_dataframes(data_dir: str, ext=".xlsx", verbose=False):
+    """
+    merge_dataframes - given a filepath, loads and attempts to merge all files as dataframes
+
+    Args:
+        data_dir (str): [root directory to search in]
+        ext (str, optional): [anticipate file extension for the dataframes ]. Defaults to '.xlsx'.
+
+    Returns:
+        pd.DataFrame(): merged dataframe of all files
+    """
+
+    src = Path(data_dir)
+    src_str = str(src.resolve())
+    mrg_df = pd.DataFrame()
+
+    all_reports = load_dir_files(directory=src_str, req_extension=ext, verbose=verbose)
+
+    failed = []
+
+    for df_path in tqdm(all_reports, total=len(all_reports), desc="joining data..."):
+
+        try:
+            this_df = pd.read_excel(df_path).convert_dtypes()
+
+            mrg_df = pd.concat([mrg_df, this_df], axis=0)
+        except Exception:
+            short_p = os.path.basename(df_path)
+            print(
+                f"WARNING - file with extension {ext} and name {short_p} could not be read."
+            )
+            failed.append(short_p)
+
+    if len(failed) > 0:
+        print("failed to merge {} files, investigate as needed")
+
+    if verbose:
+        pp.pprint(mrg_df.info(True))
+
+    return mrg_df
+
+
+def download_URL(url: str, file=None, dlpath=None, verbose=False):
+    """
+    download_URL - download a file from a URL and show progress bar
+
+    Parameters
+    ----------
+    url : str,        URL to download
+    file : str, optional, default None, name of file to save to. If None, will use the filename from the URL
+    dlpath : str, optional, default None, path to save the file to. If None, will save to the current working directory
+    verbose : bool, optional, default False, print progress bar
+
+    Returns
+    -------
+    str - path to the downloaded file
+    """
+
+    if file is None:
+        if "?dl=" in url:
+            # is a dropbox link
+            prefile = url.split("/")[-1]
+            filename = str(prefile).split("?dl=")[0]
+        else:
+            filename = url.split("/")[-1]
+
+        file = clean(filename)
+    if dlpath is None:
+        dlpath = Path.cwd()  # save to current working directory
+    else:
+        dlpath = Path(dlpath)  # make a path object
+
+    r = requests.get(url, stream=True, allow_redirects=True)
+    total_size = int(r.headers.get("content-length"))
+    initial_pos = 0
+    dl_loc = dlpath / file
+    with open(str(dl_loc.resolve()), "wb") as f:
+        with tqdm(
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            desc=file,
+            initial=initial_pos,
+            ascii=True,
+        ) as pbar:
+            for ch in r.iter_content(chunk_size=1024):
+                if ch:
+                    f.write(ch)
+                    pbar.update(len(ch))
+
+    if verbose:
+        print(f"\ndownloaded {file} to {dlpath}\n")
+
+    return str(dl_loc.resolve())
+
+
+def dl_extract_zip(
+    URLtoget: str,
+    extract_loc: str = None,
+    file_header: str = "TEMP_archive_dl_",
+    verbose: bool = False,
+):
+    """
+    dl_extract_zip - generic function to download a zip file and extract it
+
+    Parameters
+    ----------
+    URLtoget : str, zip file URL to download
+    extract_loc : str, optional, default None, path to save the zip file to. If None, will save to the current working directory
+    file_header : str, optional, default 'TEMP_archive_dl_', prefix for the zip file name
+    verbose : bool, optional, default False, print progress bar
+
+    Returns
+    -------
+    str - path to the downloaded and extracted folder
+    """
+
+    extract_loc = Path(extract_loc)
+    extract_loc.mkdir(parents=True, exist_ok=True)
+
+    save_loc = download_URL(
+        url=URLtoget, file=f"{file_header}.zip", dlpath=None, verbose=verbose
+    )
+
+    shutil.unpack_archive(save_loc, extract_dir=extract_loc)
+
+    if verbose:
+        print("extracted zip file - ", datetime.now())
+        x = load_dir_files(extract_loc, req_extension="", verbose=verbose)
+
+    # remove original
+    try:
+        os.remove(save_loc)
+        del save_loc
+    except Exception as e:
+        warnings.warn(message=f"unable to delete original zipfile due to {e}")
+    if verbose:
+        print("finished extracting zip - ", datetime.now())
+
+    return extract_loc