ein.py

from collections import UserString
from typing import Any

# Taken from
# https://www.irs.gov/businesses/small-businesses-self-employed/valid-eins
ALLOWED_PREFIXES: set[str] = {
    "01",
    "02",
    "03",
    "04",
    "05",
    "06",
    "10",
    "11",
    "12",
    "13",
    "14",
    "15",
    "16",
    "20",
    "21",
    "22",
    "23",
    "24",
    "25",
    "26",
    "27",
    "30",
    "31",
    "32",
    "33",
    "34",
    "35",
    "36",
    "37",
    "38",
    "39",
    "40",
    "41",
    "42",
    "43",
    "44",
    "45",
    "46",
    "47",
    "48",
    "50",
    "51",
    "52",
    "53",
    "54",
    "55",
    "56",
    "57",
    "58",
    "59",
    "60",
    "61",
    "62",
    "63",
    "64",
    "65",
    "66",
    "67",
    "68",
    "71",
    "72",
    "73",
    "74",
    "75",
    "76",
    "77",
    "80",
    "81",
    "82",
    "83",
    "84",
    "85",
    "86",
    "87",
    "88",
    "90",
    "91",
    "92",
    "93",
    "94",
    "95",
    "98",
    "99",
}


class EIN(UserString):
    __slots__ = ()

    def __init__(self, seq: Any) -> None:  # noqa: ANN401
        clean = str(seq).strip().replace("-", "")

        if len(clean) != 9:
            raise ValueError(f"Invalid EIN length: {len(clean)}")

        if not clean.isdigit():
            raise ValueError("Invalid EIN format")

        if clean[:2] not in ALLOWED_PREFIXES:
            raise ValueError(f"Invalid EIN prefix: {clean[:2]!r}")

        super().__init__(clean)

    def __repr__(self) -> str:
        return f'EIN("{self.data[:2]}-{self.data[2:]}")'