import pynisher


def sleepy(x: int) -> int:
    time.sleep(x)
    return x

# You can also use `cpu_time` instead
with pynisher.limit(sleepy, wall_time=7) as limited_sleep:
    x = limited_sleep(10)  # Will raise a TimeoutException

from pynisher import limit, MemoryLimitException, WallTimeoutException


def train_memory_hungry_model(X, y) -> Model:
    # ... do some thing
    return model

model_trainer = limit(
    train_memory_hungry_model,
    memory=(500, "MB"),
    wall_time=(1.5, "h")  # 1h30m
)

try:
    model = model_trainer(X, y)
except (WallTimeoutException, MemoryLimitException):
    model = None

from pynisher import limit, EMPTY

def f():
    raise ValueError()

limited_f = limit(f, wall_time=(2, "m"), raises=False)
result = limited_f()

if result is not EMPTY:
    # ...

from pynisher import restricted

@restricted(wall_time=1, raises=False)
def notify_remote_server() -> Response:
    """We don't care that this fails, just give it a second to try"""
    server = block_until_access(...)
    response = server.notify()

notify_remote_server()
# ... continue on even if it failed

from pynisher import limit


def f():
    raise ValueError()

limited_f = limit(f)

try:
    limited_f()
except ValueError as e:
    ... # do what you need

from pathlib import Path
import pickle

from pynisher import limit

def train_gpt3(save_path: Path) -> bool:
    gpt3 = ...
    gpt3.train()
    with save_path.open('wb') as f:
        pickle.dump(gpt3, f)

    return True

path = Path('gpt3.model')
trainer = limit(train_gpt3, memory=(1_000_000, "gb")):

try:
    trainer(save_path=path)

    with path.open("rb") as f:
        gpt3 = pickle.load(f)

except MemoryLimitException as e:
    ...

from pynisher import limit


for limit in ["cpu_time", "wall_time", "memory", "decorator"]:
    print(f"Supports {limit} - {supports(limit)}")


limited_f = limit(f, ...)
if not limited_f.supports("memory"):
    ...

# The name given to the multiprocessing.Process
name: str | None = None


# The memory limit to place. Specify the amount of bytes or (int, unit) where unit
# can be "B", "KB", "MB" or "GB"
memory: int | tuple[int, str] | None = None


# The cpu time in seconds to limit the process to. This time is only counted while the
# process is active.
# Can provide in (time, units) such as (1.5, "h") to indicate one and a half hours.
# Units available are "s", "m", "h"
cpu_time: int | tuple[float, str] | None = None


# The wall time in seconds to limit the process to
# Can provide in (time, units) such as (1.5, "h") to indicate one and a half hours.
# Units available are "s", "m", "h"
wall_time: int | tuple[float, str] | None = None


# Whether to throw any errors that occured in the subprocess or to silently
# throw them away. If `True` and an Error was raised, `None` will be returned.
# The errors raised in the subprocess will be the same type that are raised in
# the controlling process. The exception to this are MemoryErrors which occur
# in the subprocess, we convert these to MemoryLimitException.
raises: bool = True


# This is the multiprocess context used, please refer to their documentation
# https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
context: "fork" | "spawn" | "forkserver" | BaseContext | None = None


# Whether to emit warnings from  limit or not. The current warnings:
# * When the memory limit is lower than the starting memory of a process
# * When trying to remove the memory limit for sending back information
#   from the subprocess to the main process
warnings: bool = True


# How to handle errors. If `bool` then this decides whether or not to wrap them in
# a pynisher exception. If `list`, you can specify which errors get wrapped in a
# pynisher exception and if `dict`, you can specify what kind of errors get wrapped
# and how. See `pynisher::Pynisher::__init__` for more details on `dict`
#
# * wrap_errors={ "memory": [ImportError, (OSError, 22)], "pynisher": [ValueError] }
#
# We check that the exception is explicitly of the same type and not just a subclass.
# This is to prevent accidentally wrapping to eagerly.
wrap_errors: bool | list[Exception] | dict = False


# Whether to terminate child processes of your limited function.
# By default, pynisher will kill any subprocesses your function may spawn. If this
# is not desired behaviour, please use `daemon=True` with your spawned subprocesses
# and set `terminate_child_processes` to `False`
terminate_child_processes: bool = True

# Whether keyboard interrupts should forceably kill any subprocess or the
# pynished function. If True, it will temrinate the process tree of
# the pynished function and then reraise the KeyboardInterrupt.
forceful_keyboard_interrupt: bool = True

class PynisherException(Exception): ...
    """When a subprocess exceeds a limit"""

class TimeoutException(PynisherException): ...
    """When a subprocess exceeds a time limit (walltime or cputime)"""

class CpuTimeoutException(TimeoutException): ...
    """When a subprocess exceeds its cpu time limit"""

class WallTimeoutException(TimeoutException):
    """When a subprocess exceeds its wall time limit"""

class MemoryLimitException(PynisherException, MemoryError):
    """When a subprocess tries to allocate memory that would take it over the limit

    This also inherits from MemoryError as it is technically a MemoryError that we
    catch and convert.
    """

from contextlib import redirect_stderr

# You can always disable warnings
limited_f = limit(func, warnings=False)

# Capture warnings in a file
# Only seems to work properly on Linux
with open("stderr.txt", "w") as stderr, redirect_stderr(stderr):
    limited_f()

with open("stderr.txt", "r") as stderr:
    print(stderr.readlines())

import sys
from pynisher import PynisherException, limit

def import_sklearn() -> None:
    """Imports sklearn into a local namespace and has an sklearn object in its args"""
    from sklearn.exceptions import NotFittedError
    from sklearn.svm import SVR

    assert "sklearn" in sys.modules.keys()
    raise NotFittedError(SVR())


if __name__ == "__main__":
    # Wrapping all errors
    lf = limit(import_sklearn, wrap_errors=True)
    try:
        lf()
    except PynisherException:
        assert "sklearn" not in sys.modules.keys()

    # Wrapping only specific errors
    lf = limit(import_sklearn, wrap_errors=["NotFittedError"])
    try:
        lf()
    except PynisherException:
        assert "sklearn" not in sys.modules.keys()

    # Wrapping that error specifically as a PynisherException
    lf = limit(import_sklearn, wrap_errors={"pynisher": ["NotFittedError"]})
    try:
        lf()
    except PynisherException:
        assert "sklearn" not in sys.modules.keys()