from __future__ import annotations

import asyncio

import os

import sys

from contextlib import asynccontextmanager

from pathlib import Path

from typing import IO, TYPE_CHECKING, Any

from .exceptions import ExecutionError, PoeException

from .helpers.eventloop import run_async

if TYPE_CHECKING:

from collections.abc import AsyncIterator, Mapping, Sequence

from .config import PoeConfig

from .context import RunContext

from .io import PoeIO

from .task.base import PoeTask, TaskSpecFactory

from .ui import PoeUi

class PoeThePoet:

"""

:param cwd:

The directory that poe should take as the current working directory,

this determines where to look for a pyproject.toml file, defaults to

``Path().resolve()``

:type cwd: Path, optional

:param config:

Either a dictionary with the same schema as a pyproject.toml file, or a

`PoeConfig <https://github.com/nat-n/poethepoet/blob/main/poethepoet/config/config.py>`_

object to use as an alternative to loading config from a file.

:type config: dict | PoeConfig, optional

:param output:

A stream for the application to write its own output to, defaults to sys.stdout

:type output: IO, optional

:param poetry_env_path:

The path to the poetry virtualenv. If provided then it is used by the

`PoetryExecutor <https://github.com/nat-n/poethepoet/blob/main/poethepoet/executor/poetry.py>`_,

instead of having to execute poetry in a subprocess to determine this.

:type poetry_env_path: str, optional

:param config_name:

The name of the file to load tasks and configuration from. If not set then poe

will search for config by the following file names: pyproject.toml

poe_tasks.toml poe_tasks.yaml poe_tasks.json

:type config_name: str, optional

:param program_name:

The name of the program that is being run. This is used primarily when

outputting help messages, defaults to "poe"

:type program_name: str, optional

:param env:

Optionally provide an alternative base environment for tasks to run with.

If no mapping is provided then ``os.environ`` is used.

:type env: dict, optional

:param suppress_args:

A sequence of identifiers for global arguments that should not be displayed in

the help message.

:type suppress_args: Sequence[str], optional

"""

cwd: Path

ui: PoeUi

config: PoeConfig

_task_specs: TaskSpecFactory | None = None

def __init__(

self,

cwd: Path | str | None = None,

config: Mapping[str, Any] | PoeConfig | None = None,

output: PoeIO | IO = sys.stdout,

poetry_env_path: str | None = None,

config_name: str | None = None,

program_name: str = "poe",

env: Mapping[str, str] | None = None,

suppress_args: Sequence[str] = ("legacy_project_root",),

):

from .config import PoeConfig

from .io import PoeIO

from .ui import PoeUi

self.cwd = Path(cwd) if cwd else Path().resolve()

if self.cwd and self.cwd.is_file():

config_name = self.cwd.name

self.cwd = self.cwd.parent

self.io = (

PoeIO(

parent=output,

make_default=True,

)

if isinstance(output, PoeIO)

else PoeIO(

output=output,

error=output,

make_default=True,

)

)

if isinstance(config, PoeConfig):

self.config = config

self.config._io = self.io

else:

self.config = PoeConfig(

cwd=self.cwd, table=config, config_name=config_name, io=self.io

)

self.io.configure(baseline=self.config.verbosity)

self.ui = PoeUi(

io=self.io,

program_name=program_name,

suppress_args=suppress_args,

)

self._poetry_env_path = poetry_env_path

self._env = env if env is not None else os.environ

def __call__(self, cli_args: Sequence[str], internal: bool = False) -> int:

"""

:param cli_args:

A sequence of command line arguments to pass to poe (i.e. sys.argv[1:])

:param internal:

Indicates that this is an internal call to run poe, e.g. from a

plugin hook.

"""

self.ui.parse_args(cli_args)

if self.ui["version"]:

self.ui.print_version()

return 0

try:

return run_async(self._call(internal))

except asyncio.CancelledError:

return 1

async def _call(self, internal: bool = False) -> int:

should_display_help = self.ui["help"] != Ellipsis

try:

await self.config.load(target_path=self.ui["project_root"])

self.io.configure(baseline=self.config.verbosity)

for task_spec in self.task_specs.load_all():

task_spec.validate(self.config, self.task_specs)

except PoeException as error:

if should_display_help:

self.print_help()

return 0

self.print_help(error=error)

return 1

if should_display_help:

self.print_help()

return 0

task = self.resolve_task(internal)

if not task:

return 1

if task.has_deps():

return await self._run_task_graph(task)

return await self._run_task(task)

def modify_verbosity(self, offset: int):

"""

Set the offset by which the verbosity level will be modified in all contexts.

This is an alternative to using the `-v` and `-q` flags on the CLI.

"""

self.io.configure(offset=offset)

@property

def task_specs(self):

if not self._task_specs:

from .task.base import TaskSpecFactory

self._task_specs = TaskSpecFactory(self.config)

return self._task_specs

def resolve_task(self, allow_hidden: bool = False) -> PoeTask | None:

from .io import PoeIO

from .task.base import TaskContext

task = tuple(self.ui["task"])

if not task:

try:

self.print_help(info="No task specified.")

except PoeException as error:

self.print_help(error=error)

return None

task_name = task[0]

if task_name not in self.config.get_tasks():

self.print_help(error=PoeException(f"Unrecognized task {task_name!r}"))

return None

if task_name.startswith("_") and not allow_hidden:

self.print_help(

error=PoeException(

"Tasks prefixed with `_` cannot be executed directly"

),

)

return None

task_spec = self.task_specs.get(task_name)

task_context = TaskContext(

config=self.config,

cwd=str(task_spec.source.cwd),

specs=self.task_specs,

ui=self.ui,

io=PoeIO(

parent=self.io,

baseline_verbosity=task_spec.options.get(

"verbosity", self.io._baseline_verbosity

),

),

)

return task_spec.create_task(invocation=task, ctx=task_context)

async def _run_task(self, task: PoeTask, context: RunContext | None = None) -> int:

async with self.run_context(existing=context) as context:

try:

task_run = await task.run(context=context)

await task_run.wait(suppress_errors=False)

return task_run.return_code or 0

except ExecutionError as error:

self.ui.print_error(error=error)

return 1

except PoeException as error:

self.print_help(error=error)

return 1

async def _run_task_graph(self, task: PoeTask) -> int:

from .task.graph import TaskExecutionGraph

async with self.run_context(multistage=True) as context:

try:

graph = TaskExecutionGraph(task, context)

except PoeException as error:

self.print_help(error=error)

return 1

except ExecutionError as error:

self.ui.print_error(error=error)

return 1

plan = graph.get_execution_plan()

for stage in plan:

for stage_task in stage:

if stage_task == task:

# The final sink task gets special treatment

return await self._run_task(stage_task, context)

try:

task_run = await stage_task.run(context=context)

await task_run.wait(suppress_errors=False)

if task_run.has_failure:

raise ExecutionError(

"Task graph aborted after failed task "

f"{stage_task.name!r}"

)

except PoeException as error:

self.print_help(error=error)

return 1

except ExecutionError as error:

self.ui.print_error(error=error)

return 1

return 0

@asynccontextmanager

async def run_context(

self,

multistage: bool = False,

existing: RunContext | None = None,

) -> AsyncIterator[RunContext]:

from .context import RunContext

if existing is not None:

yield existing

return

async with RunContext.scope(

config=self.config,

ui=self.ui,

env=self._env,

dry=self.ui["dry_run"],

poe_active=self._env.get("POE_ACTIVE"),

multistage=multistage,

cwd=self.cwd,

) as context:

if self._poetry_env_path:

# This allows the PoetryExecutor to use the venv from poetry directly

# this is used by the poetry plugin

context.exec_cache["poetry_virtualenv"] = self._poetry_env_path

yield context

def print_help(

self,

info: str | None = None,

error: str | PoeException | None = None,

):

from .task.args import PoeTaskArgs

if isinstance(error, str):

error = PoeException(error)

all_tasks = self.config.get_tasks()

tasks_help: dict[

str, tuple[str, Sequence[tuple[tuple[str, ...], str, str]], str | None]

] = {

task_name: (

task.get("help", ""),

PoeTaskArgs.get_help_content(

task.get("args"), task_name, suppress_errors=bool(error)

),

task.group.name if task.group else None,

)

for task_name, task in all_tasks.items()

}

groups_headings = {

task.group.name: task.group.heading

for task in all_tasks.values()

if task.group

}

self.ui.print_help(

tasks=tasks_help, groups=groups_headings, info=info, error=error

)