library.py

"""A library of functions"""
import csv
import json
from io import StringIO

from model import Employee, Company


def read_file(filename: str) -> str:
    """Reads and returns the contents of the file."""
    with open(filename) as infile:
        return infile.read()


def parse_csv(buffer: str) -> list[dict]:
    """Parses a string of CSV int a list of dictionaries."""
    data = StringIO(buffer)
    reader = csv.DictReader(data, quotechar='"')
    return [row for row in reader]


def parse_json(buffer: str) -> list[dict]:
    """Parses a json blob into a list of dictionaries"""
    return json.loads(buffer)


def filter_employee_records(records: list[dict]) -> list[dict]:
    """Only returns records that have a employee_id field that is populated."""
    return list(filter(lambda x: bool(x.get("employee_id")), records))


def filter_company_records(records: list[dict]) -> list[dict]:
    """Only returns records that have a company_id field that is populated."""
    return list(filter(lambda x: bool(x.get("company_id")), records))


def to_pydantic(records: list[dict]) -> list[Employee | Company]:
    """Convert dictionaries to the appropriate Pydantic model"""
    return [Company(**r) if r.get("company_id") else Employee(**r) for r in records]