pcswmm_ai/layers.py

from __future__ import annotations

from typing import Any

from .exceptions import LayerNotFoundError
from .project import ProjectService


DEFAULT_LAYER_NAMES = [
    "Rain Gages", "Subcatchments", "Junctions", "Outfalls",
    "Dividers", "Storage Units", "Conduits", "Pumps",
    "Orifices", "Weirs", "Outlets",
]


class LayerService:
    def __init__(self, pcpy: Any, project: ProjectService):
        self._pcpy = pcpy
        self._project = project

    def get(self, layer_name: str):
        self._project.require_open()
        try:
            return self._pcpy.Map.Layer[layer_name]
        except Exception as exc:
            raise LayerNotFoundError(
                f"PCSWMM map layer not found: {layer_name}"
            ) from exc

    def names(self) -> list[str]:
        self._project.require_open()

        for attribute_name in ("Names", "keys", "Keys"):
            try:
                value = getattr(self._pcpy.Map.Layer, attribute_name)
                value = value() if callable(value) else value
                names = [str(item) for item in list(value)]
                if names:
                    return sorted(set(names))
            except Exception:
                pass

        names = []
        for layer_name in DEFAULT_LAYER_NAMES:
            try:
                self.get(layer_name)
                names.append(layer_name)
            except LayerNotFoundError:
                pass
        return names

    def attributes(self, layer_name: str) -> list[str]:
        layer = self.get(layer_name)
        try:
            return [str(name) for name in list(layer.Attributes)]
        except Exception:
            return []

    def entities(self, layer_name: str, query: str | None = None) -> list[Any]:
        layer = self.get(layer_name)
        if query:
            return list(layer.get_entities(query))
        return list(layer.get_entities())

    def records(self, layer_name: str, query: str | None = None) -> list[dict[str, Any]]:
        attributes = self.attributes(layer_name)
        records = []

        for entity in self.entities(layer_name, query=query):
            row = {}
            for attribute in attributes:
                try:
                    row[attribute] = entity[attribute]
                except Exception:
                    row[attribute] = None
            records.append(row)

        return records