[python]基于动态实例的命令处理设计

前言

最近在做公司内部的一个聊天机器人服务，这个聊天机器人暂时不会用到现在热门的大模型技术，只是用于接收用户固定格式的命令，然后调用对应的方法。因为只是内部使用，所以性能也不需要太高。目前考虑的用户命令类型有以下几种：

单命令。比如用户发一个ping，调用ping主命令。
有一个子命令。比如用户发送ping version，调用ping主命令的version子命令。
单命令，带一系列位置参数。比如ping host1 host2 host3，调用ping主命令，主命令自行处理参数。
子命令有一系列位置参数。比如ping tcp host1 host2 host3，调用ping主命令的tcp子命令来处理参数。

暂不考虑子命令的子命令、flag等命令形式。

早期也没想着搞太复杂的功能，所以代码用正则表达式匹配，然后写了一堆if ... else，如今看来不是很美观，而且每次新增命令都要去配置下匹配逻辑，给别人修改时，别人经常忘了改匹配逻辑，比较繁琐。

这版的修改想法是命令类一旦声明就自动注册到某个地方，接收命令的时候自动分发到对应的命令类及其方法。想到的几个方案有监听者模式、责任链模式和本文所要提的动态实例方式（我也不知道这种方法怎么命名，瞎起了个名字）。

代码结构

复制代码

│  .gitignore
│  main.py
│  README.md
│
└─commands
        cmda.py
        cmdb.py
        __init__.py

子命令的代码都存放在./commands目录下，./commands/__init__.py声明了命令的基类，导入commands目录下除了__init__.py之外的所有python文件，以及声明工厂函数。

除了__init__.py，commands目录下的所有python文件都是命令的实现。

基类

基类的声明位于commands/__init__.py文件中，要求子类必须实现main_cmd()方法，以及通过类属性判断是否需要导入命令类。自动注册子类的方法见__init_subclass__()

python 复制代码

from pathlib import Path
from abc import ABCMeta, abstractmethod
from threading import Lock
from collections import UserDict
import importlib
from functools import wraps
import inspect
from typing import Callable

class ThreadSafeDict(UserDict):
    """线程安全的字典"""
    def __init__(self):
        super().__init__()
        self._lock = Lock()
    
    def __setitem__(self, key, item):
        with self._lock:
            super().__setitem__(key, item)

class Command(metaclass=ABCMeta):
    registry = ThreadSafeDict()

    def __init__(self):
        # self._sub_cmds = ThreadSafeDict()
        self._sub_cmd: str = ""
        self._cmd_args: list = []

    @abstractmethod
    def main_cmd(self):
        pass

    @sub_cmd(name="help")
    def get_help(self):
        """Get help info"""
        message = f"Usage: {self._main_name} [subcommand] [args]\n"
        for name, f in self._sub_cmds.items():
            doc = f.__doc__ or ""
            message += f"  {name}, {doc}\n"
        print(message)

    def parse_cmd(self):
        cmd_list = self.command.split(" ")
        cmd_list_length = len(cmd_list)
        if cmd_list_length == 1:
            self._sub_cmd = ""
            self._cmd_args = []
        elif cmd_list_length >= 2 and cmd_list[1] not in self._sub_cmds:
            self._sub_cmd = ""
            self._cmd_args = cmd_list[1:]
        elif cmd_list_length >= 2 and cmd_list[1] in self._sub_cmds:
            self._sub_cmd = cmd_list[1]
            self._cmd_args = cmd_list[2:]
        else:
            self._sub_cmd = ""
            self._cmd_args = []

    def dispatch_command(self) -> Callable:
        """
        根据主命令和子命令的名称分发到相应的命令处理方法

        Returns:
            Callable: 返回对应的命令处理方法, 如果找不到匹配的子命令则返回 None
        """
        if not self._sub_cmd and not self._cmd_args:
            return self.main_cmd
        elif not self._sub_cmd and self._cmd_args:
            return self.main_cmd
        elif self._sub_cmd and self._sub_cmd not in self._sub_cmds:
            return None
        else:
            return self._sub_cmds[self._sub_cmd]
        
    def run(self):
        self.parse_cmd()
        func = self.dispatch_command()
        if not func:
            self.get_help()
        else:
            func(self)

	def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls_main_name = getattr(cls, "_main_name", "")
        cls_enabled = getattr(cls, "_enabled", False)
        cls_description = getattr(cls, "_description", "")
        if cls_main_name and cls_enabled and cls_description:
            cls.registry[cls._main_name.lower()] = cls  # 自动注册子类
            if not hasattr(cls, "_sub_cmds"):
                cls._sub_cmds = ThreadSafeDict()
            for name, method in inspect.getmembers(cls, inspect.isfunction):
                if hasattr(method, "__sub_cmd__"):
                    cls._sub_cmds[method.__sub_cmd__] = method
        else:
            print(f"{cls.__name__} 未注册，请检查类属性 _main_name, _enabled, _description")

子类只有导入时才会自动注册，所以写了个遍历目录进行导入的函数。

python 复制代码

def load_commands(dir_path: Path) -> None:
    """遍历目录下的所有python文件并导入"""
    commands_dir = Path(dir_path)
    for py_file in commands_dir.glob("*.py"):
        if py_file.stem in ("__init__"):
            continue
        module_name = f"commands.{py_file.stem}"
        try:
            importlib.import_module(module_name)
        except ImportError as e:
            print(f"Failed to import {module_name}: {e}")

load_commands(Path(__file__).parent)

子命令装饰器

命令类可以使用装饰器来注册子命令，其实只是给函数加个属性。

python 复制代码

def sub_cmd(name: str):
    """
    装饰器函数, 用于包装目标函数并添加 __sub_cmd 属性

    Args:
        name (str): 子命令名称
    """
    def decorator(func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            return func(self, *args, **kwargs)
        wrapper.__sub_cmd__ = name
        return wrapper
    return decorator

实现命令类

随便写两个命令类。命令类必须声明_main_name、_enabled和_description这三个类属性，否则不会注册这个命令类。

cmda

代码文件为commands/cmda.py

python 复制代码

from commands import Command, sub_cmd

class Cmda(Command):
    _main_name = "cmda"
    _enabled = True
    _description = "this is cmda"

    def __init__(self, command: str):
        self.command = command
        super().__init__()

    def main_cmd(self, *args: tuple, **kwargs):
        print("this is main cmd for cmda")
    
    @sub_cmd(name="info")
    def get_info(self):
        """Get info"""
        print(f"this is cmda's info")

cmdb

代码文件为commands/cmdb.py

python 复制代码

from commands import Command, sub_cmd

class Cmdb(Command):
    _main_name = "cmdb"
    _enabled = True
    _description = "this is cmdb"

    def __init__(self, command: str):
        self.command = command
        super().__init__()

    def main_cmd(self, *args, **kwargs):
        print("this is cmdb main")


    @sub_cmd("info")
    def get_info(self):
        print("this is cmdb info")
        if self._cmd_args:
            print(f"args: {self._cmd_args}")

工厂函数

工厂函数的代码也是位于commands/__init__.py

python 复制代码

def create_command(command: str) -> Command:
    """工厂函数"""
    if not command:
        raise ValueError("command can not be empty")
    command_list = command.split(" ")
    command_type = command_list[0]
    cls = Command.registry.get(command_type.lower())
    if not cls:
        raise ValueError(f"Unknown command: {command_type}")
    return cls(command)

使用示例

使用示例的代码位于main.py

python 复制代码

from commands import create_command

if __name__ == '__main__':
    command = create_command("cmdb info aaa")
    command.run()
    command = create_command("cmda help")
    command.run()

执行输出

复制代码

this is cmdb info
args: ['aaa']
Usage: cmda [subcommand] [args]
  help, Get help info
  info, Get info

完整代码

除了commands/__init__.py，其它代码文件的完整内容上面都有了，所以补充下__init__.py的内容

python 复制代码

from pathlib import Path
from abc import ABCMeta, abstractmethod
from threading import Lock
from collections import UserDict
import importlib
from functools import wraps
import inspect
from typing import Callable

def sub_cmd(name: str):
    """
    装饰器函数, 用于包装目标函数并添加 __sub_cmd 属性

    Args:
        name (str): 子命令名称
    """
    def decorator(func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            return func(self, *args, **kwargs)
        wrapper.__sub_cmd__ = name
        return wrapper
    return decorator

class ThreadSafeDict(UserDict):
    """线程安全的字典"""
    def __init__(self):
        super().__init__()
        self._lock = Lock()
    
    def __setitem__(self, key, item):
        with self._lock:
            super().__setitem__(key, item)

class Command(metaclass=ABCMeta):
    registry = ThreadSafeDict()

    def __init__(self):
        # self._sub_cmds = ThreadSafeDict()
        self._sub_cmd: str = ""
        self._cmd_args: list = []

    @abstractmethod
    def main_cmd(self):
        pass

    @sub_cmd(name="help")
    def get_help(self):
        """Get help info"""
        message = f"Usage: {self._main_name} [subcommand] [args]\n"
        for name, f in self._sub_cmds.items():
            doc = f.__doc__ or ""
            message += f"  {name}, {doc}\n"
        print(message)

    def parse_cmd(self):
        cmd_list = self.command.split(" ")
        cmd_list_length = len(cmd_list)
        if cmd_list_length == 1:
            self._sub_cmd = ""
            self._cmd_args = []
        elif cmd_list_length >= 2 and cmd_list[1] not in self._sub_cmds:
            self._sub_cmd = ""
            self._cmd_args = cmd_list[1:]
        elif cmd_list_length >= 2 and cmd_list[1] in self._sub_cmds:
            self._sub_cmd = cmd_list[1]
            self._cmd_args = cmd_list[2:]
        else:
            self._sub_cmd = ""
            self._cmd_args = []

    def dispatch_command(self) -> Callable:
        """
        根据主命令和子命令的名称分发到相应的命令处理方法

        Returns:
            Callable: 返回对应的命令处理方法, 如果找不到匹配的子命令则返回 None
        """
        if not self._sub_cmd and not self._cmd_args:
            return self.main_cmd
        elif not self._sub_cmd and self._cmd_args:
            return self.main_cmd
        elif self._sub_cmd and self._sub_cmd not in self._sub_cmds:
            return None
        else:
            return self._sub_cmds[self._sub_cmd]
        
    def run(self):
        self.parse_cmd()
        func = self.dispatch_command()
        if not func:
            self.get_help()
        else:
            func(self)



    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls_main_name = getattr(cls, "_main_name", "")
        cls_enabled = getattr(cls, "_enabled", False)
        cls_description = getattr(cls, "_description", "")
        if cls_main_name and cls_enabled and cls_description:
            cls.registry[cls._main_name.lower()] = cls  # 自动注册子类
            if not hasattr(cls, "_sub_cmds"):
                cls._sub_cmds = ThreadSafeDict()
            for name, method in inspect.getmembers(cls, inspect.isfunction):
                if hasattr(method, "__sub_cmd__"):
                    cls._sub_cmds[method.__sub_cmd__] = method
        else:
            print(f"{cls.__name__} 未注册，请检查类属性 _main_name, _enabled, _description")

def create_command(command: str) -> Command:
    """工厂函数"""
    if not command:
        raise ValueError("command can not be empty")
    command_list = command.split(" ")
    command_type = command_list[0]
    cls = Command.registry.get(command_type.lower())
    if not cls:
        raise ValueError(f"Unknown command: {command_type}")
    return cls(command)

def load_commands(dir_path: Path) -> None:
    """遍历目录下的所有python文件并导入"""
    commands_dir = Path(dir_path)
    for py_file in commands_dir.glob("*.py"):
        if py_file.stem in ("__init__"):
            continue
        module_name = f"commands.{py_file.stem}"
        try:
            importlib.import_module(module_name)
        except ImportError as e:
            print(f"Failed to import {module_name}: {e}")

load_commands(Path(__file__).parent)

__all__ = [
    "create_command",
]