Enrich an onboarding record with a company's filing and background

import json
import os
import sys
from urllib.error import HTTPError
from urllib.parse import quote
from urllib.request import Request, urlopen


class DatavilleError(Exception):
    """Raised when Dataville returns an error other than no results."""


def lookup(source, keywords):
    """Search one Dataville source and return its data object, or None if nothing matched."""
    api_key = os.environ.get("DATAVILLE_API_KEY")
    if not api_key:
        raise DatavilleError("Set DATAVILLE_API_KEY to your Dataville API key.")
    request = Request(
        f"https://api.dataville.com/{source}/{quote(keywords, safe='')}",
        headers={"Authorization": f"Bearer {api_key}"},
    )
    try:
        with urlopen(request) as response:
            return json.load(response)["data"]
    except HTTPError as error:
        # A 404 means the source has nothing for these keywords.
        if error.code == 404:
            return None
        raise DatavilleError(f"Dataville returned {error.code} {error.reason} for {source}.") from None


def enrich(business_name):
    """Return the filing and background fields to store with an application."""
    record = {"filing": None, "background": None, "stale": False}

    # Filing: when the business isn't an SEC filer, the edgar source returns
    # another company's filing that mentions the name, marked with a
    # search_query field. Keep only a result without one.
    filing_result = lookup("edgar", business_name)
    if filing_result and "search_query" not in filing_result["metadata"]:
        filing_metadata = filing_result["metadata"]
        record["filing"] = {
            "filer_name": filing_metadata["company_name"],
            "form_type": filing_metadata["form_type"],
            "filing_date": filing_metadata["filing_date"],
            "accession_number": filing_metadata["accession_number"],
            "document_url": filing_metadata["document_url"],
        }

    # Background: keep the article only if its title names the business and it
    # isn't a disambiguation page, whose abstract reads "X may refer to".
    # Wikipedia titles often drop a leading "The", so compare without it.
    article_result = lookup("wiki", business_name)
    if article_result:
        article_title = article_result["title"]
        abstract = article_result["metadata"]["abstract"] or ""
        normalized_name = business_name.lower().removeprefix("the ").strip()
        title_matches = normalized_name in article_title.lower()
        is_disambiguation = "may refer to" in abstract
        if title_matches and not is_disambiguation:
            record["background"] = {
                "title": article_title,
                "url": article_result["metadata"]["url"],
                "abstract": abstract,
                "last_updated": article_result["last_updated"],
            }

    # A stale result came from Dataville's stored copy while a source was down.
    # Flag the record so you can run enrich again later.
    record["stale"] = any("stale" in result for result in (filing_result, article_result) if result)
    return record


if __name__ == "__main__":
    # Try it from a terminal: python3 enrich.py "The Home Depot"
    if len(sys.argv) != 2:
        raise SystemExit('Usage: python3 enrich.py "Business legal name"')
    try:
        print(json.dumps(enrich(sys.argv[1]), indent=2))
    except DatavilleError as error:
        raise SystemExit(error)