Skip to content

Latest commit

 

History

History
307 lines (230 loc) · 10.4 KB

File metadata and controls

307 lines (230 loc) · 10.4 KB

Adding a New Data Source

This guide walks through the process of adding a new data source collector to the codecarbon-data project.

Overview

Each data source is a collector — a Python class that fetches raw data from an external source and normalizes it into typed Pydantic records. The pipeline then merges records from multiple collectors within the same domain and writes a validated CSV.

Architecture:

Collector.collect()          →  raw data (any format)
Collector.normalize()        →  list[DomainRecord]   (typed Pydantic models)
MergePolicy.merge()          →  list[DomainRecord]   (deduplicated/reconciled)
records_to_dataframe()       →  pd.DataFrame         (CSV boundary only)
DataFrame.to_csv()           →  data/{domain}.csv
validate_csv()               →  schema check against datapackage.json

Step-by-Step

1. Choose the domain

Every collector belongs to one domain:

Domain Output CSV Record Model
cpu hardware_cpu.csv CpuRecord
gpu hardware_gpu.csv GpuRecord
grid grid_emissions.csv GridRecord
cloud cloud_emissions.csv CloudRecord
embodied embodied_carbon.csv EmbodiedRecord

2. Register the source in sources.yaml

Add an entry under sources: with the source's licensing metadata:

sources:
  my_source:
    name: "Human-readable Source Name"
    url: "https://example.com/data"
    license: "CC BY 4.0"              # or "MIT", "US Public Domain", "Proprietary (Vendor)", etc.
    redistribution: "permissive"       # or "factual_extraction"
    access_method: "http_api"          # or "http_scrape", "http_download", "browser_scrape", "local_package"
    attribution: "Data sourced from Example (https://example.com), licensed under CC BY 4.0."
    constraints: "Optional notes about usage restrictions."  # optional

The source_name on your collector class must match the key here (my_source). The pipeline validates this linkage at startup.

3. Create the collector class

Create src/collectors/{domain}/my_source.py:

"""My source collector — collects {domain} data from Example API."""

import logging

from src.base import BaseCollector, CollectorError
from src.http import create_client
from src.models import CpuRecord  # or GpuRecord, GridRecord, CloudRecord

logger = logging.getLogger(__name__)

SOURCE_URL = "https://example.com/api/data"


class MySourceCollector(BaseCollector):
    """Collects data from the Example API."""

    domain = "cpu"            # must match a pipeline domain
    source_name = "my_source" # must match the key in sources.yaml
    requires_browser = False  # True if Playwright is needed

    def __init__(self):
        self._raw_data = None

    def collect(self) -> list[dict]:
        """Fetch raw data from the Example API."""
        client = create_client()
        try:
            response = client.get(SOURCE_URL)
            response.raise_for_status()
            self._raw_data = response.json()
        except Exception as e:
            raise CollectorError(
                self.source_name,
                "Failed to fetch data from Example API",
                str(e),
            ) from e
        finally:
            client.close()

        if not self._raw_data:
            raise CollectorError(self.source_name, "API returned empty data")

        logger.info("Loaded %d entries from Example API", len(self._raw_data))
        return self._raw_data

    def normalize(self) -> list[CpuRecord]:
        """Transform raw data to typed CpuRecord instances."""
        if self._raw_data is None:
            raise CollectorError(
                self.source_name, "collect() must be called before normalize()"
            )

        records = []
        for entry in self._raw_data:
            records.append(
                CpuRecord(
                    manufacturer=entry["vendor"],
                    model=entry["name"],
                    tdp_w=float(entry["tdp"]) if entry.get("tdp") else None,
                    # ... map remaining fields ...
                    source_url=SOURCE_URL,
                    data_source=self.source_name,
                )
            )

        records.sort(key=lambda r: (r.manufacturer, r.model))
        logger.info("Normalized %d records", len(records))
        return records

Key rules:

  • collect() fetches raw data in any format. It should raise CollectorError on failure.
  • normalize() returns a list of the domain's Pydantic record type. Never return DataFrames.
  • Use create_client() from src.http for all HTTP requests (provides retries, rate limiting, User-Agent).
  • Set data_source=self.source_name on each record.
  • Sort records deterministically for reproducible output.

4. Write tests

Create tests/collectors/{domain}/test_my_source.py. Use inline sample data and mock create_client instead of reading from fixture files. This keeps tests self-contained and avoids stale fixtures.

"""Tests for the My Source collector."""

from unittest.mock import MagicMock, patch

import pytest

from src.base import CollectorError
from src.collectors.cpu.my_source import MySourceCollector
from src.models import CpuRecord

# Inline sample data — small representative subset (3-5 entries)
SAMPLE_API_RESPONSE = [
    {"vendor": "Acme", "name": "CPU-100", "tdp": "65"},
    {"vendor": "Acme", "name": "CPU-200", "tdp": "95"},
]


class TestMySourceCollector:
    def _mock_collect(self):
        """Helper: mock HTTP and run collect()."""
        mock_response = MagicMock()
        mock_response.json.return_value = SAMPLE_API_RESPONSE
        mock_response.raise_for_status = MagicMock()

        mock_client = MagicMock()
        mock_client.get.return_value = mock_response

        with patch(
            "src.collectors.cpu.my_source.create_client",
            return_value=mock_client,
        ):
            collector = MySourceCollector()
            collector.collect()
        return collector

    def test_collect_loads_data(self):
        """Test that collect() fetches data successfully."""
        collector = self._mock_collect()
        assert collector._raw_data is not None

    def test_normalize_returns_typed_records(self):
        """Test that normalize() returns CpuRecord instances."""
        collector = self._mock_collect()
        records = collector.normalize()
        assert all(isinstance(r, CpuRecord) for r in records)

    def test_normalize_before_collect_raises(self):
        """Test that normalize() raises if collect() hasn't been called."""
        collector = MySourceCollector()
        with pytest.raises(CollectorError):
            collector.normalize()

    def test_records_have_required_fields(self):
        """Test that all records have non-empty manufacturer and model."""
        collector = self._mock_collect()
        records = collector.normalize()
        for record in records:
            assert record.manufacturer
            assert record.model
            assert record.source_url
            assert record.data_source == "my_source"

    def test_records_sorted_deterministically(self):
        """Test that output is sorted."""
        collector = self._mock_collect()
        records = collector.normalize()
        models = [r.model for r in records]
        assert models == sorted(models)

5. Register in the pipeline

Edit src/pipeline.py in _get_domain_collectors():

elif domain == "cpu":
    from src.collectors.cpu.intel_ark import IntelArkCpuCollector
    from src.collectors.cpu.my_source import MySourceCollector  # ADD

    collectors.append(IntelArkCpuCollector())
    collectors.append(MySourceCollector())  # ADD

6. Update the merge policy (if needed)

If the new source overlaps with existing sources (same entities from different sources), update the merge policy in src/collectors/{domain}/merge.py to handle deduplication.

For example, the CPU merge policy uses manufacturer-authoritative sources — Intel ARK is authoritative for Intel CPUs, AMD specs for AMD. Supplementary sources fill in gaps:

# In src/collectors/cpu/merge.py
MANUFACTURER_SOURCE = {
    "Intel": "intel_ark",
    "AMD": "amd_specs",
}

If your source adds a new manufacturer or provides supplementary data, the existing gap-filling logic handles it automatically. If it needs special treatment, modify the merge policy.

For non-overlapping sources (like different cloud providers), the merge is a simple concatenation — no changes needed.

7. Update the schema (if adding new fields)

If your source provides data for fields not yet in datapackage.json, add them:

  1. Add the field to the Pydantic model in src/models.py (as Optional[type] = None)
  2. Add the field to the schema in datapackage.json under the appropriate resource
  3. Run uv run pytest tests/test_validate.py to verify schema consistency

8. Verify

# Run all tests
uv run pytest

# Lint
uv run ruff check src/ tests/

# Run the pipeline for your domain
uv run python -m src.pipeline --domain cpu

# Inspect output
head data/hardware_cpu.csv

# Or use the CLI
uv run python -m src.cli collect --domain cpu

Testing Strategy

Collectors always fetch live data — there are no fixture fallbacks. For tests, mock the HTTP layer using unittest.mock.patch on create_client:

from unittest.mock import MagicMock, patch

SAMPLE_DATA = [{"name": "Test", "value": 42}]

mock_response = MagicMock()
mock_response.json.return_value = SAMPLE_DATA
mock_response.raise_for_status = MagicMock()

mock_client = MagicMock()
mock_client.get.return_value = mock_response

with patch(
    "src.collectors.cpu.my_source.create_client",
    return_value=mock_client,
):
    collector = MySourceCollector()
    collector.collect()
    records = collector.normalize()

This approach keeps tests fast, self-contained, and free from stale fixture files.

File Checklist

When adding a new source, you should create or modify these files:

  • sources.yaml — add licensing entry
  • src/collectors/{domain}/my_source.py — the collector class
  • tests/collectors/{domain}/test_my_source.py — unit tests (with HTTP mocks)
  • src/pipeline.py — register in _get_domain_collectors()
  • src/collectors/{domain}/merge.py — update if overlapping data (optional)
  • src/models.py — add new fields if needed (optional)
  • datapackage.json — add new schema fields if needed (optional)