"""Internal helper for parameterless modules bound to the one source layer."""

from __future__ import annotations

import weakref

from torch import nn

from .source import DendroSourceLayer


class SourceBoundModule(nn.Module):
    """An ``nn.Module`` that references, but never registers, the source module."""

    def __init__(self, source: DendroSourceLayer) -> None:
        super().__init__()
        self.__dict__["_source_ref"] = weakref.ref(source)

    @property
    def source(self) -> DendroSourceLayer:
        source = self.__dict__["_source_ref"]()
        if source is None:
            raise RuntimeError("The owning DendroSourceLayer has been released")
        return source