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

1#!/usr/bin/env python 

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

3 

4import os 

5import re 

6import sys 

7from functools import wraps 

8from importlib.resources import read_text 

9 

10from rich.panel import Panel 

11 

12from .. import T, __respkg__ 

13 

14def 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 

26def 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 

44def 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