1  from __future__ import annotations
  2  
  3  from importlib import import_module
  4  from typing import Callable, Type
  5  
  6  from ..context import Context, COMMAND_FUNCTION, COMPLETER_FUNCTION, ContextType
  7  from ..exceptions import CommandRegistrationError, NoDocStringError
  8  
  9  
 10  class CommandError(Exception):
 11      def __init__(self, msg: str):
 12          super().__init__()
 13          self.msg: str = msg
 14  
 15  
 16  class Command:
 17      def __init__(
 18          self,
 19          name: str,
 20          func: COMMAND_FUNCTION[ContextType],
 21          description: str,
 22          contexts: list[Type[Context]],
 23          aliases: list[str],
 24      ):
 25          self.name: str = name
 26          self.func: COMMAND_FUNCTION[ContextType] = func
 27          self.description: str = description
 28          self.contexts: list[Type[Context]] = contexts
 29          self.aliases: list[str] = aliases
 30          self.completer_func: COMPLETER_FUNCTION[ContextType] | None = None
 31          self.subcommands: list[Command] = []
 32          self.prepared_subcommands: dict[Type[Context], dict[str, Command]] = {}
 33  
 34      def __call__(self, context: ContextType, args: list[str]) -> None:
 35          self.func(context, args)
 36  
 37      def parse_command(self, context: Context, args: list[str]) -> tuple[Command, list[str]]:
 38          prepared_subcommands = self.prepared_subcommands.get(type(context), {})
 39          if args:
 40              cmd: Command | None = prepared_subcommands.get(args[0])
 41              if cmd is not None:
 42                  return cmd.parse_command(context, args[1:])
 43  
 44          return self, args
 45  
 46      def handle(self, context: ContextType, args: list[str]) -> None:
 47          cmd, args = self.parse_command(context, args)
 48          try:
 49              cmd.func(context, args)
 50          except CommandError as error:
 51              print(error.msg)
 52  
 53      def handle_completer(self, context: ContextType, args: list[str]) -> list[str]:
 54          cmd, args = self.parse_command(context, args)
 55  
 56          out = list(cmd.prepared_subcommands.get(type(context), {}))
 57          if cmd.completer_func is not None:
 58              out += cmd.completer_func(context, args) or []
 59  
 60          return out
 61  
 62      def completer(self) -> Callable[[COMPLETER_FUNCTION[ContextType]], COMPLETER_FUNCTION[ContextType]]:
 63          def decorator(func: COMPLETER_FUNCTION[ContextType]) -> COMPLETER_FUNCTION[ContextType]:
 64              self.completer_func = func
 65              return func
 66  
 67          return decorator
 68  
 69      def subcommand(
 70          self, name: str, *, contexts: list[Type[Context]] | None = None, aliases: list[str] | None = None
 71      ) -> Callable[[COMMAND_FUNCTION[ContextType]], "Command"]:
 72          def decorator(func: COMMAND_FUNCTION[ContextType]) -> Command:
 73              if func.__doc__ is None:
 74                  raise NoDocStringError(name, subcommand=True)
 75              desc: str = func.__doc__
 76              desc = "\n".join(map(str.strip, desc.splitlines())).strip()
 77              cmd = Command(name, func, desc, contexts or self.contexts, aliases or [])
 78              self.subcommands.append(cmd)
 79              return cmd
 80  
 81          return decorator
 82  
 83      def make_subcommands(self) -> None:
 84          for cmd in self.subcommands:
 85              for context in cmd.contexts:
 86                  for name in [cmd.name] + cmd.aliases:
 87                      if name in self.prepared_subcommands.setdefault(context, {}):
 88                          raise CommandRegistrationError(name, subcommand=True)
 89                      self.prepared_subcommands[context][name] = cmd
 90              cmd.make_subcommands()
 91  
 92  
 93  commands: list[Command] = []
 94  
 95  
 96  def command(
 97      name: str, contexts: list[Type[Context]], aliases: list[str] | None = None
 98  ) -> Callable[[COMMAND_FUNCTION[ContextType]], Command]:
 99      def decorator(func: COMMAND_FUNCTION[ContextType]) -> Command:
100          if func.__doc__ is None:
101              raise NoDocStringError(name)
102          desc: str = func.__doc__
103          desc = "\n".join(map(str.strip, desc.splitlines())).strip()
104          cmd = Command(name, func, desc, contexts, aliases or [])
105          commands.append(cmd)
106          return cmd
107  
108      return decorator
109  
110  
111  def make_commands() -> dict[Type[Context], dict[str, Command]]:
112      for module in [
113          "account",
114          "help",
115          "shell",
116          "status",
117          "device",
118          "files",
119          "morphcoin",
120          "service",
121          "miner",
122          "inventory",
123          "shop",
124          "network",
125      ]:
126          import_module(f"PyCrypCli.commands.{module}")
127  
128      result: dict[Type[Context], dict[str, Command]] = {}
129      for cmd in commands:
130          for context in cmd.contexts:
131              for name in [cmd.name] + cmd.aliases:
132                  if name in result.setdefault(context, {}):
133                      raise CommandRegistrationError(name)
134                  result[context][name] = cmd
135          cmd.make_subcommands()
136      return result