diff options
Diffstat (limited to 'utilfn.py')
-rw-r--r-- | utilfn.py | 122 |
1 files changed, 122 insertions, 0 deletions
diff --git a/utilfn.py b/utilfn.py new file mode 100644 index 0000000..e062d1b --- /dev/null +++ b/utilfn.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# A set of functions and types to be included for general use. +from enum import Enum +from typing import * +import time +import modin.pandas as pd +from dataclasses import dataclass + +# Local imports +from utiltypes import * + + +# Returns a displayable time, use e.g. with print(f"{delta_now(time_begin)}") +def delta_now(t: float, dp: int = 2) -> float: + return round(time.time() - t, dp) + +def ask_file(prompt: str = "What file should be opened: ") -> str: + filename = input(prompt) + return filename + +def open_file(file: str) -> pd.DataFrame: + if file.endswith(".csv"): + return pd.read_csv(file) + else: + return pd.read_excel(file) + +# Returns None if successfully saved, else a str error message. +# TODO: add error handling. +def save_df(df: pd.DataFrame, filetype: Filetype, name: str, sheet_name = None): + if sheet_name == None: + sheet_name = name + + if filetype == Filetype.Excel: + df.to_excel(f"{name}.xlsx", sheet_name = f"{sheet_name}", index = False) + else: + df.to_csv(f"{name}.csv", index = False) + +@dataclass +class PrintOptions: + update_user: bool + show_timings: bool + +def print_noendl(s: str): + print(s, end = '', flush = True) + +def reprint_noendl(s: str): + print(f'\r{s}', end = '', flush = True) + +# Please supply your own \n if you need it. +def cond_print(print_options: PrintOptions, s: str, timing_str: str = None, suffix = ""): + if timing_str and print_options.show_timings: + print(s, timing_str, end = "", flush = True) + elif print_options.update_user: + print(s, end = "", flush = True) + print(suffix, end = "", flush = True) + +@dataclass +class AskOpenFileReply: + df: pd.DataFrame + fname: str + sheet: str + time_to_open_file: float + +# Asks for a filename, and opens it; keeps opening until a valid file passed. +def ask_open_file(show_timings: bool = True) -> AskOpenFileReply: + while True: + try: + file = ask_file() + if file.lower().strip() in [ "quit", "exit" ]: + print("Quitting...") + quit() + + print("Opening file... ", end="", flush=True) + # TODO: add in to open a specific sheet; prompt if more than one + # sheet. At present it just opens the first one. Use optparam + # sheet_name=, with either str of sheet name or int index (deflt 0) + start_time = time.time() + df = open_file(file) + time_taken = time.time() - start_time + if show_timings: print(f" opened in {round(time_taken, 2)}s!") + else: print( " opened!") + break + except Exception as exc: + print(" File failed to open or could not be found.", flush=True) + print(f"Details: {exc}") + + # TODO: make `sheet` sheet value when I add this logic. + reply = AskOpenFileReply(df = df, fname = file, sheet = "", time_to_open_file = time_taken) + return reply + +def ask_excel_or_csv(given: str = None, prompt: bool = True) -> Filetype: + # prompt: if False, throw if given not valid instead of prompting the user. + exit_options = [ "exit", "quit" ] + excel_options = [ "xl", "excel", "xlsx", ".xlsx" ] + csv_options = [ "csv", ".csv" ] + + if given: + given = given.lower().strip() + if given in exit_options: + quit() + elif given in excel_options: + return Filetype.Excel + elif given in csv_options: + return Filetype.Csv + # Do not add an else clause as we need to handle the case when `given` + # is `None` in the below prompt. + + if not prompt: + # TODO: change this to be a more specific exception type. + raise Exception(f"ask_excel_or_csv: passed {given} as parameter with prompt == False") + + answer = input("Should I output as `.xlsx` or `.csv` (`csv` is much faster)? [csv, excel, quit] ") + answer = answer.lower().strip() + while True: + if answer in exit_options: + quit() + elif answer in csv_options: + return Filetype.Csv + elif answer in excel_options: + return Filetype.Excel + else: + print("Invalid input - I want one of `csv`, `excel`, or `quit`.") |