Coverage for aipyapp/aipy/utils.py: 23%

1#!/usr/bin/env python

2# -*- coding: utf-8 -*-

3

4 import os

5 import re

6 import sys

7 from functools import wraps

8 from importlib.resources import read_text

9

10 from rich.panel import Panel

11

12 from .. import T, __respkg__

13

14 def restore_output(func):

15 @wraps(func)

16 def wrapper(self, *args, **kwargs):

17 old_stdout, old_stderr = sys.stdout, sys.stderr

18 sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__

19

20 try:

21 return func(self, *args, **kwargs)

22 finally:

23 sys.stdout, sys.stderr = old_stdout, old_stderr

24 return wrapper

25

26 def confirm_disclaimer(console):

27 DISCLAIMER_TEXT = read_text(__respkg__, "DISCLAIMER.md")

28 console.print()

29 panel = Panel.fit(DISCLAIMER_TEXT, title="[red]免责声明", border_style="red", padding=(1, 2))

30 console.print(panel)

31

32 while True:

33 console.print("\n[red]是否确认已阅读并接受以上免责声明？[/red](yes/no):", end=" ")

34 response = input().strip().lower()

35 if response in ("yes", "y"):

36 console.print("[green]感谢确认，程序继续运行。[/green]")

37 return True

38 elif response in ("no", "n"):

39 console.print("[red]您未接受免责声明，程序将退出。[/red]")

40 return False

41 else:

42 console.print("[yellow]请输入 yes 或 no。[/yellow]")

43

44 def get_safe_filename(input_str, extension=".html", max_length=16):

45 input_str = input_str.strip()

46 safe_str = re.sub(r'[\\/:*?"<>|\s]', '', input_str).strip()

47 if not safe_str:

48 return None

49

50 name = safe_str[:max_length]

51 base_name = name

52 filename = f"{base_name}{extension}" if extension else base_name

53 counter = 1

54

55 while os.path.exists(filename):

56 filename = f"{base_name}_{counter}{extension}" if extension else f"{base_name}_{counter}"

57 counter += 1

58

59 return filename