# Path: xmanager/xm/async_packager.py
# class PackageHasNotBeenCalledError(RuntimeError):
# class AsyncPackager:
#     class PicklablePackageImpl:
# async def _return_executable(
#     executable: job_blocks.Executable) -> job_blocks.Executable:
#   def __init__(
#       self, package_batch: Callable[[Sequence[job_blocks.Packageable]],
#                                     Sequence[job_blocks.Executable]]
#   ) -> None:
#   def add(
#       self,
#       packageable: job_blocks.Packageable) -> Awaitable[job_blocks.Executable]:
#     def check_is_packaged() -> None:
#     async def package_impl() -> job_blocks.Executable:
#       def __await__(self):
#       def __reduce__(self):
#   def package(
#       self, extra_packageables: Sequence[job_blocks.Packageable] = ()
#   ) -> Sequence[job_blocks.Executable]:
# 
# Path: xmanager/xm/job_blocks.py
# class SequentialArgs:
#   class _RegularItem:
#   class _KeywordItem:
# class ExecutableSpec(abc.ABC):
# class Executable(abc.ABC):
# class ExecutorSpec(abc.ABC):
# class Executor(abc.ABC):
# class Packageable:
# class Constraint(abc.ABC):
# class Job:
# class JobGroup:
#   def __init__(self) -> None:
#   def _ingest_regular_item(self, value: Any) -> None:
#   def _ingest_keyword_item(self, name: str, value: Any) -> None:
#   def _merge_from(self, args: 'SequentialArgs') -> None:
#     def import_regular_item(item: SequentialArgs._RegularItem):
#     def import_keyword_item(item: SequentialArgs._KeywordItem):
#   def from_collection(collection: Optional[UserArgs]) -> 'SequentialArgs':
#     def check_for_string(args: str) -> None:
#     def import_sequential_args(args: SequentialArgs) -> None:
#     def import_mapping(collection: Mapping[Any, Any]) -> None:
#     def import_sequence(collection: Sequence[Any]) -> None:
#   def to_list(
#       self,
#       escaper: Callable[[Any], str],
#       kwargs_joiner: Callable[[str, str], str] = utils.trivial_kwargs_joiner
#   ) -> List[str]:
#     def export_regular_item(item: SequentialArgs._RegularItem) -> str:
#     def export_keyword_item(item: SequentialArgs._KeywordItem) -> str:
#   def to_dict(self, kwargs_only: bool = False) -> Dict[str, Any]:
#     def export_regular_item(
#         item: SequentialArgs._RegularItem) -> Tuple[str, Any]:
#     def export_keyword_item(
#         item: SequentialArgs._KeywordItem) -> Tuple[str, Any]:
#   def __eq__(self, other) -> bool:
#   def __repr__(self) -> str:
# def merge_args(*operands: Union[SequentialArgs, UserArgs]) -> SequentialArgs:
#   def name(self) -> str:
#   def Spec(cls) -> ExecutorSpec:  # pylint: disable=invalid-name
#   def __init__(self,
#                *,
#                constraints: Optional[Sequence[Constraint]] = None,
#                **jobs: JobType) -> None:
# 
# Path: xmanager/xm/utils.py
# FLAGS = flags.FLAGS
# ARG_ESCAPER = pattern_matching.match(
#     pattern_matching.Case([ShellSafeArg], lambda v: v.arg),
#     pattern_matching.Case([enum.Enum], lambda v: shlex.quote(str(v.name))),
#     pattern_matching.Case([Any], lambda v: shlex.quote(str(v))),
# )
# class SpecialArg(abc.ABC):
# class ShellSafeArg(SpecialArg):
# def trivial_kwargs_joiner(key: str, value: str) -> str:
# def run_in_asyncio_loop(
#     f: Callable[..., ReturnType]) -> Callable[..., ReturnType]:
#   def decorated(*args, **kwargs) -> ReturnType:
# def find_launch_script_path() -> str:
# def resolve_path_relative_to_launcher(path: str) -> str:


import pickle
import unittest
from typing import Sequence
from xmanager.xm import async_packager
from xmanager.xm import job_blocks
from xmanager.xm import utils
  ]


class _TestExecutableSpec(job_blocks.ExecutableSpec):

  def __init__(self, name: str) -> None:
    self._name = name

  @property
  def name(self) -> str:
    return self._name


def _make_packageable(name: str) -> job_blocks.Packageable:
  return job_blocks.Packageable(
      executable_spec=_TestExecutableSpec(name),
      executor_spec=job_blocks.ExecutorSpec())


class AsyncPackagerTest(unittest.TestCase):

  @utils.run_in_asyncio_loop
  async def test_async_packager_end_to_end(self):
    packager = async_packager.AsyncPackager(_package_batch)
    executable1 = packager.add(_make_packageable('1'))
    executable2 = packager.add(_make_packageable('2'))
    packager.package()
    self.assertEqual((await executable1).name, '1')
    self.assertEqual((await executable2).name, '2')


# Path: xmanager/xm/async_packager.py
# class PackageHasNotBeenCalledError(RuntimeError):
# class AsyncPackager:
#     class PicklablePackageImpl:
# async def _return_executable(
#     executable: job_blocks.Executable) -> job_blocks.Executable:
#   def __init__(
#       self, package_batch: Callable[[Sequence[job_blocks.Packageable]],
#                                     Sequence[job_blocks.Executable]]
#   ) -> None:
#   def add(
#       self,
#       packageable: job_blocks.Packageable) -> Awaitable[job_blocks.Executable]:
#     def check_is_packaged() -> None:
#     async def package_impl() -> job_blocks.Executable:
#       def __await__(self):
#       def __reduce__(self):
#   def package(
#       self, extra_packageables: Sequence[job_blocks.Packageable] = ()
#   ) -> Sequence[job_blocks.Executable]:
# 
# Path: xmanager/xm/job_blocks.py
# class SequentialArgs:
#   class _RegularItem:
#   class _KeywordItem:
# class ExecutableSpec(abc.ABC):
# class Executable(abc.ABC):
# class ExecutorSpec(abc.ABC):
# class Executor(abc.ABC):
# class Packageable:
# class Constraint(abc.ABC):
# class Job:
# class JobGroup:
#   def __init__(self) -> None:
#   def _ingest_regular_item(self, value: Any) -> None:
#   def _ingest_keyword_item(self, name: str, value: Any) -> None:
#   def _merge_from(self, args: 'SequentialArgs') -> None:
#     def import_regular_item(item: SequentialArgs._RegularItem):
#     def import_keyword_item(item: SequentialArgs._KeywordItem):
#   def from_collection(collection: Optional[UserArgs]) -> 'SequentialArgs':
#     def check_for_string(args: str) -> None:
#     def import_sequential_args(args: SequentialArgs) -> None:
#     def import_mapping(collection: Mapping[Any, Any]) -> None:
#     def import_sequence(collection: Sequence[Any]) -> None:
#   def to_list(
#       self,
#       escaper: Callable[[Any], str],
#       kwargs_joiner: Callable[[str, str], str] = utils.trivial_kwargs_joiner
#   ) -> List[str]:
#     def export_regular_item(item: SequentialArgs._RegularItem) -> str:
#     def export_keyword_item(item: SequentialArgs._KeywordItem) -> str:
#   def to_dict(self, kwargs_only: bool = False) -> Dict[str, Any]:
#     def export_regular_item(
#         item: SequentialArgs._RegularItem) -> Tuple[str, Any]:
#     def export_keyword_item(
#         item: SequentialArgs._KeywordItem) -> Tuple[str, Any]:
#   def __eq__(self, other) -> bool:
#   def __repr__(self) -> str:
# def merge_args(*operands: Union[SequentialArgs, UserArgs]) -> SequentialArgs:
#   def name(self) -> str:
#   def Spec(cls) -> ExecutorSpec:  # pylint: disable=invalid-name
#   def __init__(self,
#                *,
#                constraints: Optional[Sequence[Constraint]] = None,
#                **jobs: JobType) -> None:
# 
# Path: xmanager/xm/utils.py
# FLAGS = flags.FLAGS
# ARG_ESCAPER = pattern_matching.match(
#     pattern_matching.Case([ShellSafeArg], lambda v: v.arg),
#     pattern_matching.Case([enum.Enum], lambda v: shlex.quote(str(v.name))),
#     pattern_matching.Case([Any], lambda v: shlex.quote(str(v))),
# )
# class SpecialArg(abc.ABC):
# class ShellSafeArg(SpecialArg):
# def trivial_kwargs_joiner(key: str, value: str) -> str:
# def run_in_asyncio_loop(
#     f: Callable[..., ReturnType]) -> Callable[..., ReturnType]:
#   def decorated(*args, **kwargs) -> ReturnType:
# def find_launch_script_path() -> str:
# def resolve_path_relative_to_launcher(path: str) -> str:


# Path: xmanager/xm/async_packager_test.py
import pickle
import unittest
from typing import Sequence
from xmanager.xm import async_packager
from xmanager.xm import job_blocks
from xmanager.xm import utils

  ]


class _TestExecutableSpec(job_blocks.ExecutableSpec):

  def __init__(self, name: str) -> None:
    self._name = name

  @property
  def name(self) -> str:
    return self._name


def _make_packageable(name: str) -> job_blocks.Packageable:
  return job_blocks.Packageable(
      executable_spec=_TestExecutableSpec(name),
      executor_spec=job_blocks.ExecutorSpec())


class AsyncPackagerTest(unittest.TestCase):

  @utils.run_in_asyncio_loop
  async def test_async_packager_end_to_end(self):
    packager = async_packager.AsyncPackager(_package_batch)
    executable1 = packager.add(_make_packageable('1'))
    executable2 = packager.add(_make_packageable('2'))
    packager.package()
    self.assertEqual((await executable1).name, '1')
    self.assertEqual((await executable2).name, '2')


  @utils.run_in_asyncio_loop