taotoolkit 1.0.0__py3-none-any.whl

taotoolkit/tao_common.py

@@ -0,0 +1,435 @@
+import os
+import re
+import subprocess
+import sys
+import time
+from functools import wraps
+from datetime import datetime
+from typing import Callable, Dict, List, Optional, Any
+from colorama import Back, Fore, Style, init
+
+init(autoreset=True)
+
+_DEBUG_MODE = False
+
+
+def debugCtrl(debug: int) -> None:
+    global _DEBUG_MODE
+    _DEBUG_MODE = debug == 1
+
+level_colors = {
+    "DEBUG": Fore.WHITE,
+    "INFO": Fore.GREEN,
+    "HIGH": Fore.MAGENTA,
+    "WARN": Fore.YELLOW,
+    "ERROR": Fore.RED,
+}
+
+def logDecorator(level="INFO", callStack=1):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                frame = sys._getframe(callStack)
+                fileName = os.path.basename(frame.f_code.co_filename)
+                funcName = frame.f_code.co_name
+                lineNum = frame.f_lineno
+                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+
+                result = func(*args, **kwargs)
+                if level == "DEBUG" and _DEBUG_MODE != 1:
+                    return result
+                color = level_colors.get(level, Fore.WHITE)
+                message = f"{timestamp} {color}{Back.RESET}{Style.NORMAL}{level.ljust(5)} [{fileName}:{lineNum} {funcName}] "
+                args_str = " ".join(map(str, args))
+                kwargs_str = " ".join(f"{k}={v}" for k, v in kwargs.items())
+                all_args = f"{args_str} {kwargs_str}".strip()
+
+                print(f"{message}: {all_args}")
+                return result
+            except Exception as e:
+                print(f"Error in logDecorator: {e}")
+                return None
+
+        return wrapper
+
+    return decorator
+
+
+@logDecorator(level="DEBUG")
+def logd(*args):
+    pass
+
+
+@logDecorator(level="INFO")
+def logi(*args):
+    pass
+
+
+@logDecorator(level="HIGH")
+def logh(*args):
+    pass
+
+
+@logDecorator(level="WARN")
+def logw(*args):
+    pass
+
+
+@logDecorator(level="ERROR")
+def loge(*args):
+    pass
+
+
+def _create_log_call_func(level, default_callStack=2):
+    """Factory function to create log functions with customizable callStack."""
+    def log_func(*args, callStack=None):
+        if callStack is None:
+            callStack = default_callStack
+        
+        try:
+            frame = sys._getframe(callStack)
+            fileName = os.path.basename(frame.f_code.co_filename)
+            funcName = frame.f_code.co_name
+            lineNum = frame.f_lineno
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+
+            if level == "DEBUG" and _DEBUG_MODE != 1:
+                return None
+            
+            color = level_colors.get(level, Fore.WHITE)
+            message = f"{timestamp} {color}{Back.RESET}{Style.NORMAL}{level.ljust(5)} [{fileName}:{lineNum} {funcName}] "
+            args_str = " ".join(map(str, args))
+
+            print(f"{message}: {args_str}")
+            return None
+        except Exception as e:
+            print(f"Error in log call: {e}")
+            return None
+    
+    return log_func
+
+
+logCallDebug = _create_log_call_func("DEBUG", default_callStack=2)
+logCallInfo = _create_log_call_func("INFO", default_callStack=2)
+logCallWarn = _create_log_call_func("WARN", default_callStack=2)
+logCallHigh = _create_log_call_func("HIGH", default_callStack=2)
+logCallErr = _create_log_call_func("ERROR", default_callStack=2)
+
+
+def exit() -> None:
+    sys.exit(1)
+
+
+def eExit(msg: str) -> None:
+    loge(msg)
+    sys.exit(1)
+
+
+def iExit(msg: str) -> None:
+    logi(msg)
+    sys.exit(0)
+
+
+def sleep(seconds: float) -> None:
+    time.sleep(seconds)
+
+
+sysVersion = sys.platform
+
+
+def createFolder(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def fileExist(path: str) -> bool:
+    return os.path.exists(path)
+
+
+def findFilesBySuffixKeyword(path: str, suffix: str, keyword: str) -> List[str]:
+    result = []
+    for root, _dirs, files in os.walk(path):
+        for f in files:
+            if f.endswith(suffix) and keyword in f:
+                result.append(os.path.join(root, f))
+    return result
+
+
+def findFilesWithSuffix(path: str, suffix: str) -> List[str]:
+    result = []
+    for root, _dirs, files in os.walk(path):
+        for f in files:
+            if f.endswith(suffix):
+                result.append(os.path.join(root, f))
+    return result
+
+
+def get_file_index(prefix: str, suffix: str, filename: str) -> int:
+    pattern = re.compile(rf"{prefix}(\d+){suffix}")
+    res = pattern.findall(filename)
+    if len(res) != 1:
+        raise Exception("Cannot get file index!")
+    return int(res[0])
+
+
+def is_number(s: str) -> bool:
+    try:
+        float(s)
+        return True
+    except ValueError:
+        return False
+
+
+def run_as_admin() -> bool:
+    try:
+        import ctypes
+
+        return ctypes.windll.shell32.IsUserAnAdmin()
+    except Exception:
+        return False
+
+
+def removeAnyContainStr(string: str, subStringList: list, matchCase=True) -> List[str]:
+    if matchCase:
+        return [s for s in subStringList if string not in s]
+    return [s for s in subStringList if string.lower() not in s.lower()]
+
+
+def renameFile(old_path: str, new_path: str) -> None:
+    os.rename(old_path, new_path)
+
+
+class ErrCode:
+    SUCCESS = 0
+    ERROR = -1
+    FILE_NOT_FOUND = -2
+    FILE_ILLEGAL = -3
+
+
+def fileExceptionNotFound(path: str) -> Exception:
+    return FileNotFoundError(f"File not found: {path}")
+
+
+def fileExceptionillegal(path: str) -> Exception:
+    return ValueError(f"Illegal file: {path}")
+
+
+def normalize_args(args: List[str]) -> List[str]:
+    if not args:
+        return []
+    first = args[0]
+    if first.endswith(".py") or first.endswith(".exe") or first in {"tao", "python", "python3"}:
+        return args[1:]
+    if os.path.basename(first) in {"tao", "tao.exe"}:
+        return args[1:]
+    return list(args)
+
+
+def register_command(*names: str, registry: Optional[Dict[str, Callable]] = None) -> Callable:
+    def decorator(func: Callable) -> Callable:
+        if registry is not None:
+            for name in names:
+                registry[name] = func
+        return func
+
+    return decorator
+
+
+def inputSerial(start: int, end: int) -> int:
+    while True:
+        try:
+            idx = int(input(f"请输入序号 ({start}-{end}): "))
+            if start <= idx <= end:
+                return idx
+            loge(f"输入超出范围，请输入 {start}-{end} 之间的数字")
+        except ValueError:
+            loge("请输入有效的数字")
+
+
+def refreshWait(_keyword: str, timeout: int, interval: float, func: Callable, *args) -> bool:
+    start = time.time()
+    while time.time() - start < timeout:
+        if func(*args):
+            return True
+        sleep(interval)
+    return False
+
+
+def checkMsg(msg: str, *keywords: str) -> None:
+    if msg:
+        for kw in keywords:
+            if kw in msg.lower():
+                eExit(msg)
+
+
+def parseArgs(argsList: List[str], prog: str, description: str = "") -> Any:
+    import argparse
+
+    parser = argparse.ArgumentParser(prog=prog, description=description)
+    return parser.parse_args(argsList)
+
+
+def parseDeviceInfo(info: str) -> Dict[str, str]:
+    result = {}
+    lines = info.strip().split("\n")
+    for line in lines:
+        if "device" in line:
+            parts = line.split()
+            if len(parts) >= 2:
+                result["serial"] = parts[0]
+                result["model"] = parts[1]
+    return result
+
+
+def pCmd(
+    cmd: str,
+    sleep_s=1.0,
+    showCmd=False,
+    shell=False,
+    splitCmd=True,
+    encoding="utf-8",
+    waitDone=True,
+    timeOut=3,
+    showRet=False,
+    checkErr=False,
+    ignoreTimeOut=False,
+    logLevel="debug",
+):
+    if showCmd:
+        if logLevel == "warn":
+            logCallWarn(cmd)
+        elif logLevel == "error":
+            logCallErr(cmd)
+        elif logLevel == "info":
+            logCallInfo(cmd)
+        else:
+            logCallDebug(cmd)
+    try:
+        if splitCmd:
+            proc = subprocess.Popen(
+                cmd.split(" "),
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                encoding=encoding,
+                shell=shell,
+                # env=os.environ.copy(),
+            )
+        else:
+            proc = subprocess.Popen(
+                cmd,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                encoding=encoding,
+                shell=shell,
+                # env=os.environ.copy(),
+            )
+    except subprocess.TimeoutExpired:
+        return "", f"Command '{cmd}' timed out after {timeOut} seconds"
+    except subprocess.CalledProcessError as e:
+        return "", f"Command '{cmd}' failed with error: {e.stderr}"
+    out = ""
+    error = ""
+    if waitDone:
+        try:
+            out, error = proc.communicate(timeout=timeOut)
+        except subprocess.TimeoutExpired:
+            logCallErr(f"{cmd} timeout", "error")
+            if not ignoreTimeOut:
+                exit()
+            else:
+                return "None", "None"
+
+        sleep(sleep_s)
+        if checkErr:
+            checkMsg(error, "error")
+            checkMsg(out, "failed")
+            checkMsg(out, "error")
+        if showRet:
+            if out != "":
+                logCallInfo(re.sub(r"\n$", "", out))
+            if error != "" and ("error" in error or "failed" in error):
+                logCallErr(error)
+            elif error != "":
+                logCallInfo(error)
+        return out, error
+    else:
+        return "None", "None"
+
+
+def rCmd(
+    cmd: str,
+    sleep_s=1.0,
+    showCmd=False,
+    shell=False,
+    splitCmd=True,
+    encoding="utf-8",
+    timeOut=3,
+    showRet=False,
+    checkErr=False,
+    logLevel="debug",
+):
+    if showCmd:
+        if logLevel == "warn":
+            logCallWarn(cmd,callStack=3)
+        elif logLevel == "error":
+            logCallErr(cmd,callStack=3)
+        elif logLevel == "info":
+            logCallInfo(cmd,callStack=3)
+        else:
+            logCallDebug(cmd,callStack=3)
+    try:
+        if splitCmd:
+            result = subprocess.run(
+                cmd.split(" "),
+                capture_output=True,
+                text=True,
+                encoding=encoding,
+                shell=shell,
+                check=True,
+                timeout=timeOut,
+            )
+        else:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                encoding=encoding,
+                shell=shell,
+                check=True,
+                timeout=timeOut,
+            )
+    except subprocess.TimeoutExpired:
+        return "", f"Command '{cmd}' timed out after {timeOut} seconds"
+    except subprocess.CalledProcessError as e:
+        return "", f"Command '{cmd}' failed with error: {e.stderr}"
+
+    sleep(sleep_s)
+    if checkErr:
+        checkMsg(result.stderr, "error")
+        checkMsg(result.stderr, "failed")
+        checkMsg(result.stdout, "error")
+    if showRet:
+        if len(result.stdout) != 0:
+            logCallInfo(result.stdout.strip())
+        if len(result.stderr) != 0:
+            if "error" in result.stderr or "failed" in result.stderr:
+                logCallErr(result.stderr.strip())
+            else:
+                logCallInfo(result.stderr.strip())
+    return result.stdout, result.stderr
+
+
+def oCmd(cmd: str, sleep_s=1.0, showCmd=False, logLevel="debug"):
+    if showCmd:
+        if logLevel == "warn":
+            logCallWarn(cmd, callStack=3)
+        elif logLevel == "error":
+            logCallErr(cmd, callStack=3)
+        elif logLevel == "info":
+            logCallInfo(cmd, callStack=3)
+        else:
+            logCallDebug(cmd, callStack=3)
+    os.system(cmd)
+    sleep(sleep_s)