main.dang

pub description = "A toolchain for testing Python applications with automatic OpenTelemetry tracing"

enum Runner {
  AUTO
  UV
  PIP
}

enum EnvTarget {
  AUTO
  VENV
  SYSTEM
}

type Pytest {
  """
  The uv binary extracted from the bundled Docker image
  """
  let uvBinary: Directory! {
    currentModule.source.directory("images/uv").dockerBuild.rootfs
  }

  """
  The alpine base image
  """
  let alpine: Container! {
    currentModule.source.directory("images/alpine").dockerBuild
  }

  """
  The pytest_otel source directory bundled with this module
  """
  let pytestOtelSource = currentModule.source.directory("pytest_otel")

  """
  The source directory containing your Python project, taken from the
  workspace at `sourcePath` in the constructor.
  """
  pub source: Directory!

  """
  Additional arguments to pass to pytest (e.g., ["-x", "--tb=short"])
  """
  pub args: [String!]! = ["-v"]

  """
  Optional: Python version to use (e.g., "3.14", "3.13", "3.12")

  Only relevant when using the default base.
  """
  pub pythonVersion: String!

  """
  Optional: a custom image that replaces the default Alpine+uv base.
  When set, Python provisioning (`uv venv`) is skipped.
  """
  pub container: Container = null

  """
  Which Python tool to use. AUTO probes the container (prefers uv, else pip).
  """
  pub runner: Runner = Runner.AUTO

  """
  Skip installing project dependencies and mounting the source code directory
  before running `test`. Set true when the container already has both baked in.
  """
  pub skipInstallDeps: Boolean!

  """
  Construct the toolchain from a workspace.

  `source` is taken from the workspace at `sourcePath`; every other field is
  exposed here so consumers and `customizations` can configure the full
  pipeline (pytest args, Python version, custom container, runner, dependency
  installation).
  """
new(
    ws: Workspace!,
    sourcePath: String! = "/",
    args: [String!]! = ["-v"],
    """
    Optional: Python version to use (e.g., "3.14", "3.13", "3.12")

    Only relevant when using the default base.
    """
    pythonVersion: String! = "3.14",
    container: Container = null,
    runner: Runner = Runner.AUTO,
    """
    Skip installing project dependencies and mounting the source code directory
    before running `test`. Set true when the container already has both baked in.
    """
    skipInstallDeps: Boolean! = false,
  ) {
    self.source = ws.directory(sourcePath)
    self.args = args
    self.pythonVersion = pythonVersion
    self.container = container
    self.runner = runner
    self.skipInstallDeps = skipInstallDeps
    self
  }

  """
  The default Alpine + uv base: Python tooling only.
  """
  let defaultBase: Container! {
    alpine
      .withExec(["apk", "add", "libgcc"])
      .withEnvVariable("PYTHONUNBUFFERED", "1")
      .withEnvVariable("PATH", "/root/.local/bin:/usr/local/bin:$PATH", expand: true)
      .withEnvVariable("UV_LINK_MODE", "copy")
      .withEnvVariable("UV_PROJECT_ENVIRONMENT", "/opt/venv")
      .withEnvVariable("VIRTUAL_ENV", "/opt/venv")
      .withMountedCache("/root/.cache/uv", cacheVolume("pytest-toolchain-uv"))
      .withDirectory("/usr/local/bin", uvBinary, include: ["uv*"])
      .withExec(["uv", "venv", "/opt/venv", "-p", pythonVersion])
  }

  """
  The base used by `test`: the custom container if provided, else the default.
  """
  let base: Container! {
    if (container != null) {
      container
    } else {
      defaultBase
    }
  }

  """
  Test a Python project with pytest and OpenTelemetry tracing.

  Builds a container from the base (custom or default), adds your source,
  optionally installs dependencies, then dispatches to the uv or pip run path
  (auto-detected unless `runner` is set explicitly).
  """
  pub test: Void @check {
    let chosen = resolveRunner(base)
    let result = if (chosen == Runner.UV) {
      testUv(installDepsUv(base))
    } else {
      testPip(installDepsPip(base))
    }
    result.sync
    null
  }

  """
  Resolve the effective runner: when `runner` is AUTO, probe the container
  (prefer uv, else pip); otherwise use the configured value.
  """
  let resolveRunner(ctr: Container!): Runner {
    if (runner == Runner.AUTO) {
      if (ctr.withExec(["sh", "-c", "command -v uv"], expect: ReturnType.ANY).exitCode == 0) {
        Runner.UV
      } else {
        Runner.PIP
      }
    } else {
      runner
    }
  }

  let installDepsUv(ctr: Container!): Container! {
    if (!skipInstallDeps) {
      # requirements.txt-only projects need explicit install; pyproject.toml
      # projects are synced by `uv run` at test time.
      let ctr = addSourceDir(ctr, source)
      if (source.exists("requirements.txt") and !source.exists("pyproject.toml")) {
        ctr = ctr.withExec(["uv", "pip", "install", "-r", "requirements.txt"])
      }
      ctr
    } else {
      ctr
    }
  }

  let installDepsPip(ctr: Container!): Container! {
    if (!skipInstallDeps) {
      let ctr = addSourceDir(ctr, source)
      if (source.exists("requirements.txt")) {
        ctr = ctr.withExec(["pip", "install", "-r", "requirements.txt"])
      } else if (source.exists("pyproject.toml")) {
        ctr = ctr.withExec(["pip", "install", "."])
      }
      ctr
    } else {
      ctr
    }
  }

  let addSourceDir(
    ctr: Container!,
    dir: Directory!,
  ): Container! {
    ctr
      .withDirectory("/app", dir)
      .withWorkdir("/app")
  }

  """
  Build the pip command that installs pytest_otel for the given `envTarget`.

  `pip` is the invocation prefix (`pip`, `python3 -m pip`, ...). AUTO tries a
  plain install and retries with `--break-system-packages` (works in a venv, a
  plain system, or an externally-managed system); VENV installs as-is and fails
  loudly when no usable env is active; SYSTEM forces `--break-system-packages`.
  """
  let pipInstall(pip: String!, envTarget: EnvTarget): String! {
    if (envTarget == EnvTarget.SYSTEM) {
      pip + " install --break-system-packages /opt/pytest_otel"
    } else if (envTarget == EnvTarget.VENV) {
      pip + " install /opt/pytest_otel"
    } else {
      pip + " install /opt/pytest_otel || " + pip + " install --break-system-packages /opt/pytest_otel"
    }
  }

  """
  Run pytest with uv on a container that already has source + deps. Installs
  pytest_otel (forced uv, into the project venv since `uv run` is venv-based),
  then `uv run pytest`. Returns the post-run container.
  """
  pub testUv(
    ctr: Container!,
  ): Container! {
    installPytestOtel(ctr, runner: Runner.UV, envTarget: EnvTarget.VENV)
      .withExec(["uv", "run", "pytest"] + args)
  }

  """
  Run pytest with pip on a container that already has source + deps. Installs
  pytest_otel (forced pip), then `python -m pytest`. Returns the post-run container.
  """
  pub testPip(
    ctr: Container!,
  ): Container! {
    installPytestOtel(ctr, runner: Runner.PIP)
      .withExec(["python", "-m", "pytest"] + args)
  }

  """
  Install pytest_otel into a container you manage yourself.

  `runner` AUTO detects the tool at runtime (uv -> pip -> python -m pip); UV or
  PIP forces that tool even when others are present.

  `envTarget` chooses where pytest_otel lands. AUTO installs into an existing
  virtual environment when one is present and otherwise falls back to a system
  install; VENV forces a virtual environment, creating it when missing; SYSTEM
  forces a system install (`uv ... --system --break-system-packages` /
  `pip ... --break-system-packages`). Use SYSTEM (or AUTO) for containers
  without a `.venv`.

  Returns the container with pytest_otel installed; run pytest yourself on it.
  """
  pub installPytestOtel(
    ctr: Container!,
    runner: Runner = Runner.AUTO,
    envTarget: EnvTarget = EnvTarget.AUTO,
  ): Container! {
    # uv resolves its `pip install` target from VIRTUAL_ENV / an active `.venv`
    # (ignoring UV_PROJECT_ENVIRONMENT), so resolve the env path explicitly into
    # $T. AUTO installs into that env when it exists, else falls back to a system
    # install (`--system --break-system-packages` is the only form uv honors for
    # a non-virtual environment; UV_BREAK_SYSTEM_PACKAGES alone is a no-op).
    # VENV creates the env when missing; SYSTEM always targets the base Python.
    let uvResolve = "T=\"${VIRTUAL_ENV:-${UV_PROJECT_ENVIRONMENT:-.venv}}\"; "
    let uvVenv = "uv pip install --python \"$T/bin/python\" /opt/pytest_otel"
    let uvSystem = "uv pip install --system --break-system-packages /opt/pytest_otel"
    let uvInstall = if (envTarget == EnvTarget.SYSTEM) {
      uvSystem
    } else if (envTarget == EnvTarget.VENV) {
      uvResolve + "[ -x \"$T/bin/python\" ] || uv venv \"$T\"; " + uvVenv
    } else {
      uvResolve + "if [ -x \"$T/bin/python\" ]; then " + uvVenv + "; else " + uvSystem + "; fi"
    }
    let installCmd = if (runner == Runner.UV) {
      uvInstall
    } else if (runner == Runner.PIP) {
      pipInstall("pip", envTarget)
    } else {
      "if command -v uv >/dev/null 2>&1; then " + uvInstall + "; "
        + "elif command -v pip >/dev/null 2>&1; then " + pipInstall("pip", envTarget) + "; "
        + "elif command -v python3 >/dev/null 2>&1; then " + pipInstall("python3 -m pip", envTarget) + "; "
        + "elif command -v python >/dev/null 2>&1; then " + pipInstall("python -m pip", envTarget) + "; "
        + "else echo 'pytest toolchain: could not find uv, pip, or python to install pytest_otel' >&2; exit 1; fi"
    }
    ctr
      .withDirectory("/opt/pytest_otel", pytestOtelSource)
      .withExec(["sh", "-c", "set -e; " + installCmd])
  }

  """
  The bundled pytest_otel library, as a Directory, for fully-manual integration
  (the package is not yet published to PyPI).
  """
  pub pytestOtel: Directory! {
    pytestOtelSource
  }
}