matchbox-db 0.2.2__py3-none-any.whl

matchbox/__init__.py

@@ -0,0 +1,13 @@
+"""Matchbox."""
+
+from matchbox.common.exceptions import MatchboxClientSettingsException
+from matchbox.common.logging import logger
+
+try:
+    # Environment variables must be loaded first for other imports to work
+    from matchbox.client import *  # noqa: E402, F403
+except MatchboxClientSettingsException:
+    logger.warning(
+        "Impossible to initialise client. "
+        "Please ignore if running in server mode. Otherwise, check your .env file.",
+    )

matchbox/client/__init__.py

@@ -0,0 +1,9 @@
+"""All client-side functionalities of Matchbox."""
+
+from matchbox.client.helpers.cleaner import process
+from matchbox.client.helpers.index import index
+from matchbox.client.helpers.selector import match, query
+from matchbox.client.models.models import make_model
+from matchbox.client.visualisation import draw_resolution_graph
+
+__all__ = ("process", "index", "match", "query", "make_model", "draw_resolution_graph")

matchbox/client/_handler.py

@@ -0,0 +1,302 @@
+"""Functions abstracting the interaction with the server API."""
+
+import time
+from collections.abc import Iterable
+from io import BytesIO
+
+import httpx
+from pyarrow import Table
+from pyarrow.parquet import read_table
+
+from matchbox.client._settings import ClientSettings, settings
+from matchbox.common.arrow import SCHEMA_MB_IDS, table_to_buffer
+from matchbox.common.dtos import (
+    BackendRetrievableType,
+    ModelAncestor,
+    ModelMetadata,
+    ModelOperationStatus,
+    NotFoundError,
+    UploadStatus,
+)
+from matchbox.common.exceptions import (
+    MatchboxClientFileError,
+    MatchboxDeletionNotConfirmed,
+    MatchboxResolutionNotFoundError,
+    MatchboxServerFileError,
+    MatchboxSourceNotFoundError,
+    MatchboxUnhandledServerResponse,
+    MatchboxUnparsedClientRequest,
+)
+from matchbox.common.graph import ResolutionGraph
+from matchbox.common.hash import hash_to_base64
+from matchbox.common.sources import Match, Source, SourceAddress
+
+URLEncodeHandledType = str | int | float | bytes
+
+
+def encode_param_value(
+    v: URLEncodeHandledType | Iterable[URLEncodeHandledType],
+) -> str | list[str]:
+    if isinstance(v, str):
+        return v
+    elif isinstance(v, (int, float)):
+        return str(v)
+    elif isinstance(v, bytes):
+        return hash_to_base64(v)
+    # Needs to be at the end, so we don't apply it to e.g. strings
+    if isinstance(v, Iterable):
+        return [encode_param_value(item) for item in v]
+    raise ValueError(f"It was not possible to parse {v} as an URL parameter")
+
+
+def url_params(
+    params: dict[str, URLEncodeHandledType | Iterable[URLEncodeHandledType]],
+) -> dict[str, str | list[str]]:
+    """Prepares a dictionary of parameters to be encoded in a URL."""
+    non_null = {k: v for k, v in params.items() if v}
+    return {k: encode_param_value(v) for k, v in non_null.items()}
+
+
+def handle_http_code(res: httpx.Response) -> httpx.Response:
+    """Handle HTTP status codes and raise appropriate exceptions."""
+    res.read()
+
+    if 299 >= res.status_code >= 200:
+        return res
+
+    if res.status_code == 400:
+        if UploadStatus.model_validate_json(res.content, strict=False):
+            error = UploadStatus.model_validate(res.json())
+            raise MatchboxServerFileError(error.details)
+        else:
+            raise RuntimeError(f"Unexpected 400 error: {res.content}")
+
+    if res.status_code == 404:
+        error = NotFoundError.model_validate(res.json())
+        if error.entity == BackendRetrievableType.SOURCE:
+            raise MatchboxSourceNotFoundError(error.details)
+        if error.entity == BackendRetrievableType.RESOLUTION:
+            raise MatchboxResolutionNotFoundError(error.details)
+        else:
+            raise RuntimeError(f"Unexpected 404 error: {error.details}")
+
+    if res.status_code == 409:
+        error = ModelOperationStatus.model_validate(res.json())
+        raise MatchboxDeletionNotConfirmed(message=error.details)
+
+    if res.status_code == 422:
+        raise MatchboxUnparsedClientRequest(res.content)
+
+    raise MatchboxUnhandledServerResponse(res.content)
+
+
+def create_client(settings: ClientSettings) -> httpx.Client:
+    """Create an HTTPX client with proper configuration."""
+    return httpx.Client(
+        base_url=settings.api_root,
+        timeout=settings.timeout,
+        event_hooks={"response": [handle_http_code]},
+    )
+
+
+CLIENT = create_client(settings=settings)
+
+# Retrieval
+
+
+def query(
+    source_address: SourceAddress,
+    resolution_name: str | None = None,
+    threshold: int | None = None,
+    limit: int | None = None,
+) -> BytesIO:
+    res = CLIENT.get(
+        "/query",
+        params=url_params(
+            {
+                "full_name": source_address.full_name,
+                # Converted to b64 by `url_params()`
+                "warehouse_hash_b64": source_address.warehouse_hash,
+                "resolution_name": resolution_name,
+                "threshold": threshold,
+                "limit": limit,
+            }
+        ),
+    )
+
+    buffer = BytesIO(res.content)
+    table = read_table(buffer)
+
+    if not table.schema.equals(SCHEMA_MB_IDS):
+        raise MatchboxClientFileError(
+            message=(
+                f"Schema mismatch. Expected:\n{SCHEMA_MB_IDS}\nGot:\n{table.schema}"
+            )
+        )
+
+    return table
+
+
+def match(
+    targets: list[SourceAddress],
+    source: SourceAddress,
+    source_pk: str,
+    resolution_name: str,
+    threshold: int | None = None,
+) -> Match:
+    target_full_names = [t.full_name for t in targets]
+    target_warehouse_hashes = [t.warehouse_hash for t in targets]
+
+    res = CLIENT.get(
+        "/match",
+        params=url_params(
+            {
+                "target_full_names": target_full_names,
+                # Converted to b64 by `url_params()`
+                "target_warehouse_hashes_b64": target_warehouse_hashes,
+                "source_full_name": source.full_name,
+                # Converted to b64 by `url_params()`
+                "source_warehouse_hash_b64": source.warehouse_hash,
+                "source_pk": source_pk,
+                "resolution_name": resolution_name,
+                "threshold": threshold,
+            }
+        ),
+    )
+
+    return [Match.model_validate(m) for m in res.json()]
+
+
+# Data management
+
+
+def index(source: Source, data_hashes: Table) -> UploadStatus:
+    """Index a Source in Matchbox."""
+    buffer = table_to_buffer(table=data_hashes)
+
+    # Upload metadata
+    metadata_res = CLIENT.post("/sources", json=source.model_dump())
+
+    upload = UploadStatus.model_validate(metadata_res.json())
+
+    # Upload data
+    upload_res = CLIENT.post(
+        f"/upload/{upload.id}",
+        files={"file": (f"{upload.id}.parquet", buffer, "application/octet-stream")},
+    )
+
+    # Poll until complete with retry/timeout configuration
+    status = UploadStatus.model_validate(upload_res.json())
+    while status.status not in ["complete", "failed"]:
+        status_res = CLIENT.get(f"/upload/{upload.id}/status")
+        status = UploadStatus.model_validate(status_res.json())
+
+        if status.status == "failed":
+            raise MatchboxServerFileError(status.details)
+
+        time.sleep(settings.retry_delay)
+
+    return status
+
+
+def get_source(address: SourceAddress) -> Source:
+    warehouse_hash_b64 = hash_to_base64(address.warehouse_hash)
+    res = CLIENT.get(f"/sources/{warehouse_hash_b64}/{address.full_name}")
+
+    return Source.model_validate(res.json())
+
+
+def get_resolution_graph() -> ResolutionGraph:
+    """Get the resolution graph from Matchbox."""
+    res = CLIENT.get("/report/resolutions")
+    return ResolutionGraph.model_validate(res.json())
+
+
+# Model management
+
+
+def insert_model(model: ModelMetadata) -> ModelOperationStatus:
+    """Insert a model in Matchbox."""
+    res = CLIENT.post("/models", json=model.model_dump())
+    return ModelOperationStatus.model_validate(res.json())
+
+
+def get_model(name: str) -> ModelMetadata:
+    res = CLIENT.get(f"/models/{name}")
+    return ModelMetadata.model_validate(res.json())
+
+
+def add_model_results(name: str, results: Table) -> UploadStatus:
+    """Upload model results in Matchbox."""
+    buffer = table_to_buffer(table=results)
+
+    # Initialise upload
+    metadata_res = CLIENT.post(f"/models/{name}/results")
+
+    upload = UploadStatus.model_validate(metadata_res.json())
+
+    # Upload data
+    upload_res = CLIENT.post(
+        f"/upload/{upload.id}",
+        files={"file": (f"{upload.id}.parquet", buffer, "application/octet-stream")},
+    )
+
+    # Poll until complete with retry/timeout configuration
+    status = UploadStatus.model_validate(upload_res.json())
+    while status.status not in ["complete", "failed"]:
+        status_res = CLIENT.get(f"/upload/{upload.id}/status")
+        status = UploadStatus.model_validate(status_res.json())
+
+        if status.status == "failed":
+            raise MatchboxServerFileError(status.details)
+
+        time.sleep(settings.retry_delay)
+
+    return status
+
+
+def get_model_results(name: str) -> Table:
+    """Get model results from Matchbox."""
+    res = CLIENT.get(f"/models/{name}/results")
+    buffer = BytesIO(res.content)
+    return read_table(buffer)
+
+
+def set_model_truth(name: str, truth: float) -> ModelOperationStatus:
+    """Set the truth threshold for a model in Matchbox."""
+    res = CLIENT.patch(f"/models/{name}/truth", json=truth)
+    return ModelOperationStatus.model_validate(res.json())
+
+
+def get_model_truth(name: str) -> float:
+    """Get the truth threshold for a model in Matchbox."""
+    res = CLIENT.get(f"/models/{name}/truth")
+    return res.json()
+
+
+def get_model_ancestors(name: str) -> list[ModelAncestor]:
+    """Get the ancestors of a model in Matchbox."""
+    res = CLIENT.get(f"/models/{name}/ancestors")
+    return [ModelAncestor.model_validate(m) for m in res.json()]
+
+
+def set_model_ancestors_cache(
+    name: str, ancestors: list[ModelAncestor]
+) -> ModelOperationStatus:
+    """Set the ancestors cache for a model in Matchbox."""
+    res = CLIENT.post(
+        f"/models/{name}/ancestors_cache", json=[a.model_dump() for a in ancestors]
+    )
+    return ModelOperationStatus.model_validate(res.json())
+
+
+def get_model_ancestors_cache(name: str) -> list[ModelAncestor]:
+    """Get the ancestors cache for a model in Matchbox."""
+    res = CLIENT.get(f"/models/{name}/ancestors_cache")
+    return [ModelAncestor.model_validate(m) for m in res.json()]
+
+
+def delete_model(name: str, certain: bool = False) -> ModelOperationStatus:
+    """Delete a model in Matchbox."""
+    res = CLIENT.delete(f"/models/{name}", params={"certain": certain})
+    return ModelOperationStatus.model_validate(res.json())

matchbox/client/_settings.py

@@ -0,0 +1,26 @@
+"""Module to load client settings from env file."""
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from matchbox.common.exceptions import MatchboxClientSettingsException
+
+
+class ClientSettings(BaseSettings):
+    api_root: str
+    timeout: float | None = None
+    retry_delay: int = 5
+    default_warehouse: str | None = None
+
+    model_config = SettingsConfigDict(
+        extra="ignore",
+        env_prefix="MB__CLIENT__",
+        env_nested_delimiter="__",
+        env_file=".env",
+        env_file_encoding="utf-8",
+    )
+
+
+try:
+    settings = ClientSettings()
+except ValueError as e:
+    raise MatchboxClientSettingsException from e

matchbox/client/clean/__init__.py

@@ -0,0 +1,29 @@
+"""Library of default cleaning functions."""
+
+from matchbox.client.clean.lib import (
+    company_name,
+    company_number,
+    drop,
+    extract_cdms_number_to_new,
+    extract_company_number_to_new,
+    extract_duns_number_to_new,
+    postcode,
+    postcode_to_area,
+)
+from matchbox.client.clean.utils import alias, cleaning_function, unnest_renest
+
+__all__ = (
+    # Cleaning functions
+    "company_name",
+    "company_number",
+    "drop",
+    "extract_cdms_number_to_new",
+    "extract_company_number_to_new",
+    "extract_duns_number_to_new",
+    "postcode",
+    "postcode_to_area",
+    # Utility functions
+    "alias",
+    "cleaning_function",
+    "unnest_renest",
+)

matchbox/client/clean/lib.py

@@ -0,0 +1,191 @@
+"""Implementation of default cleaning functions."""
+
+from functools import partial
+
+from pandas import DataFrame
+
+from matchbox.client.clean import steps
+from matchbox.client.clean import utils as cu
+
+
+def company_name(
+    df: DataFrame,
+    column: str,
+    column_secondary: str = None,
+    stopwords: str = cu.STOPWORDS,
+) -> DataFrame:
+    """Standard cleaning function for company names.
+
+    * Lower case, remove punctuation & tokenise the company name into an array
+    * Extract tokens into: 'unusual' and 'stopwords'. Dedupe. Sort alphabetically
+    * Untokenise the unusual words back to a string
+
+    Args:
+        df: a dataframe
+        column: a column containing the company's main name
+        column_secondary: a column containing an array of the company's
+            secondary names
+        stopwords: a list of stopwords to use for this clean
+
+    Returns:
+        dataframe: the same as went in, but cleaned
+    """
+    remove_stopwords = partial(steps.remove_stopwords, stopwords=stopwords)
+
+    clean_primary = cu.cleaning_function(
+        steps.clean_punctuation,
+        steps.expand_abbreviations,
+        steps.tokenise,  # returns array
+        remove_stopwords,
+        steps.list_join_to_string,  # returns col
+    )
+
+    clean_secondary = cu.unnest_renest(clean_primary)
+
+    df = clean_primary(df, column)
+
+    if column_secondary is not None:
+        df = clean_secondary(df, column_secondary)
+
+    return df
+
+
+def company_number(df: DataFrame, column: str) -> DataFrame:
+    """Remove non-numbers, and then leading zeroes.
+
+    Args:
+        df: a dataframe
+        column: a column containing a company number
+
+    Returns:
+        dataframe: the same as went in, but cleaned
+    """
+    clean_number = cu.cleaning_function(steps.remove_notnumbers_leadingzeroes)
+
+    df = clean_number(df, column)
+
+    return df
+
+
+def postcode(df: DataFrame, column: str) -> DataFrame:
+    """Removes all punctuation, converts to upper, removes all spaces.
+
+    Args:
+        df: a dataframe
+        column: a column containing a postcode
+
+    Returns:
+        dataframe: the same as went in, but cleaned
+
+    """
+    clean_postcode = cu.cleaning_function(
+        steps.punctuation_to_spaces, steps.to_upper, steps.remove_whitespace
+    )
+
+    df = clean_postcode(df, column)
+
+    return df
+
+
+def postcode_to_area(df: DataFrame, column: str) -> DataFrame:
+    """Extracts postcode area from a postcode.
+
+    Args:
+        df: a dataframe
+        column: a column containing a postcode
+
+    Returns:
+        dataframe: the same as went in, but cleaned
+    """
+    extract_area = cu.cleaning_function(steps.get_postcode_area)
+
+    df = extract_area(df, column)
+
+    return df
+
+
+def extract_company_number_to_new(
+    df: DataFrame, column: str, new_column: str
+) -> DataFrame:
+    """Detects the Companies House CRN in a column and moves it to a new column.
+
+    Args:
+        df: a dataframe
+        column: a column containing some company numbers
+        new_column: the name of the column to add
+
+    Returns:
+        dataframe: the same as went in with a new column for CRNs
+    """
+    clean_crn = cu.cleaning_function(
+        steps.clean_punctuation_except_hyphens,
+        steps.to_upper,
+        steps.filter_company_number,
+    )
+
+    clean_crn_aliased = cu.alias(clean_crn, alias=new_column)
+
+    df = clean_crn_aliased(df, column)
+
+    return df
+
+
+def extract_duns_number_to_new(
+    df: DataFrame, column: str, new_column: str
+) -> DataFrame:
+    """Detects the Dun & Bradstreet DUNS nuber in a column and moves it to a new column.
+
+    Args:
+        df: a dataframe
+        column: a column containing some DUNS numbers
+        new_column: the name of the column to add
+
+    Returns:
+        dataframe: the same as went in with a new column for DUNs numbers
+    """
+    clean_duns = cu.cleaning_function(
+        steps.clean_punctuation_except_hyphens, steps.to_upper, steps.filter_duns_number
+    )
+
+    clean_duns_aliased = cu.alias(clean_duns, alias=new_column)
+
+    df = clean_duns_aliased(df, column)
+
+    return df
+
+
+def extract_cdms_number_to_new(
+    df: DataFrame, column: str, new_column: str
+) -> DataFrame:
+    """Detects the CDMS nuber in a column and moves it to a new column.
+
+    Args:
+        df: a dataframe
+        column: a column containing some CDMS numbers
+        new_column: the name of the column to add
+
+    Returns:
+        dataframe: the same as went in with a new column for CDMS numbers
+    """
+    clean_cdms = cu.cleaning_function(
+        steps.clean_punctuation_except_hyphens, steps.to_upper, steps.filter_cdms_number
+    )
+
+    clean_cdms_aliased = cu.alias(clean_cdms, alias=new_column)
+
+    df = clean_cdms_aliased(df, column)
+
+    return df
+
+
+def drop(df: DataFrame, column: str) -> DataFrame:
+    """Drops the column from the dataframe.
+
+    Args:
+        df: a dataframe
+        column: a column
+
+    Returns:
+        dataframe: the same as went in without the column
+    """
+    return df.drop(columns=[column])

matchbox/client/clean/steps/__init__.py

@@ -0,0 +1,71 @@
+"""Low-level components of default cleaning functions."""
+
+from matchbox.client.clean.steps.clean_basic import (
+    array_except,
+    array_intersect,
+    clean_punctuation,
+    clean_punctuation_except_hyphens,
+    dedupe_and_sort,
+    expand_abbreviations,
+    filter_cdms_number,
+    filter_company_number,
+    filter_duns_number,
+    get_digits_only,
+    get_low_freq_char_sig,
+    get_postcode_area,
+    list_join_to_string,
+    periods_to_nothing,
+    punctuation_to_spaces,
+    regex_extract_list_of_strings,
+    regex_remove_list_of_strings,
+    remove_notnumbers_leadingzeroes,
+    remove_stopwords,
+    remove_whitespace,
+    to_lower,
+    to_upper,
+    tokenise,
+    trim,
+)
+from matchbox.client.clean.steps.clean_basic_original import (
+    cms_original_clean_cdms_id,
+    cms_original_clean_ch_id,
+    cms_original_clean_company_name_ch,
+    cms_original_clean_company_name_general,
+    cms_original_clean_email,
+    cms_original_clean_postcode,
+)
+
+__all__ = (
+    # Basic steps
+    "array_except",
+    "array_intersect",
+    "periods_to_nothing",
+    "punctuation_to_spaces",
+    "clean_punctuation",
+    "clean_punctuation_except_hyphens",
+    "dedupe_and_sort",
+    "expand_abbreviations",
+    "filter_cdms_number",
+    "filter_company_number",
+    "filter_duns_number",
+    "get_digits_only",
+    "get_low_freq_char_sig",
+    "get_postcode_area",
+    "list_join_to_string",
+    "regex_extract_list_of_strings",
+    "regex_remove_list_of_strings",
+    "remove_notnumbers_leadingzeroes",
+    "remove_stopwords",
+    "remove_whitespace",
+    "to_lower",
+    "to_upper",
+    "tokenise",
+    "trim",
+    # Original CMS steps
+    "cms_original_clean_cdms_id",
+    "cms_original_clean_ch_id",
+    "cms_original_clean_company_name_ch",
+    "cms_original_clean_company_name_general",
+    "cms_original_clean_email",
+    "cms_original_clean_postcode",
+)