neofetch_util.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286
  1. from __future__ import annotations
  2. import inspect
  3. import os
  4. import platform
  5. import re
  6. import shlex
  7. import subprocess
  8. import zipfile
  9. from dataclasses import dataclass
  10. from pathlib import Path
  11. from subprocess import check_output
  12. from tempfile import TemporaryDirectory
  13. from urllib.request import urlretrieve
  14. import pkg_resources
  15. from typing_extensions import Literal
  16. from hyfetch.color_util import color
  17. from .constants import GLOBAL_CFG, MINGIT_URL
  18. from .presets import ColorProfile
  19. from .serializer import from_dict
  20. RE_NEOFETCH_COLOR = re.compile('\\${c[0-9]}')
  21. def term_size() -> tuple[int, int]:
  22. """
  23. Get terminal size
  24. :return:
  25. """
  26. try:
  27. return os.get_terminal_size().columns, os.get_terminal_size().lines
  28. except Exception:
  29. return 100, 20
  30. def ascii_size(asc: str) -> tuple[int, int]:
  31. """
  32. Get distro ascii width, height ignoring color code
  33. :param asc: Distro ascii
  34. :return: Width, Height
  35. """
  36. return max(len(line) for line in re.sub(RE_NEOFETCH_COLOR, '', asc).split('\n')), len(asc.split('\n'))
  37. def normalize_ascii(asc: str) -> str:
  38. """
  39. Make sure every line are the same width
  40. """
  41. w = ascii_size(asc)[0]
  42. return '\n'.join(line + ' ' * (w - ascii_size(line)[0]) for line in asc.split('\n'))
  43. def fill_starting(asc: str) -> str:
  44. """
  45. Fill the missing starting placeholders.
  46. E.g. "${c1}...\n..." -> "${c1}...\n${c1}..."
  47. """
  48. new = []
  49. last = ''
  50. for line in asc.split('\n'):
  51. new.append(last + line)
  52. # Line has color placeholders
  53. matches = RE_NEOFETCH_COLOR.findall(line)
  54. if len(matches) > 0:
  55. # Get the last placeholder for the next line
  56. last = matches[-1]
  57. return '\n'.join(new)
  58. @dataclass
  59. class ColorAlignment:
  60. mode: Literal['horizontal', 'vertical', 'custom']
  61. # custom_colors[ascii color index] = unique color index in preset
  62. custom_colors: dict[int, int] = ()
  63. # Foreground/background ascii color index
  64. fore_back: tuple[int, int] = ()
  65. @classmethod
  66. def from_dict(cls, d: dict):
  67. return from_dict(cls, d)
  68. def recolor_ascii(self, asc: str, preset: ColorProfile) -> str:
  69. """
  70. Use the color alignment to recolor an ascii art
  71. :return Colored ascii, Uncolored lines
  72. """
  73. asc = fill_starting(asc)
  74. if self.fore_back and self.mode in ['horizontal', 'vertical']:
  75. fore, back = self.fore_back
  76. # Replace foreground colors
  77. asc = asc.replace(f'${{c{fore}}}', color('&0' if GLOBAL_CFG.is_light else '&f'))
  78. lines = asc.split('\n')
  79. # Add new colors
  80. if self.mode == 'horizontal':
  81. colors = preset.with_length(len(lines))
  82. asc = '\n'.join([l.replace(f'${{c{back}}}', colors[i].to_ansi()) + color('&r') for i, l in enumerate(lines)])
  83. else:
  84. raise NotImplementedError()
  85. # Remove existing colors
  86. asc = re.sub(RE_NEOFETCH_COLOR, '', asc)
  87. elif self.mode in ['horizontal', 'vertical']:
  88. # Remove existing colors
  89. asc = re.sub(RE_NEOFETCH_COLOR, '', asc)
  90. lines = asc.split('\n')
  91. # Add new colors
  92. if self.mode == 'horizontal':
  93. colors = preset.with_length(len(lines))
  94. asc = '\n'.join([colors[i].to_ansi() + l + color('&r') for i, l in enumerate(lines)])
  95. else:
  96. asc = '\n'.join(preset.color_text(line) + color('&r') for line in lines)
  97. else:
  98. preset = preset.unique_colors()
  99. # Apply colors
  100. color_map = {ai: preset.colors[pi].to_ansi() for ai, pi in self.custom_colors.items()}
  101. for ascii_i, c in color_map.items():
  102. asc = asc.replace(f'${{c{ascii_i}}}', c)
  103. return asc
  104. def get_command_path() -> str:
  105. """
  106. Get the absolute path of the neofetch command
  107. :return: Command path
  108. """
  109. return pkg_resources.resource_filename(__name__, 'scripts/neowofetch')
  110. def ensure_git_bash() -> Path:
  111. """
  112. Ensure git bash installation for windows
  113. :returns git bash path
  114. """
  115. if platform.system() == 'Windows':
  116. # Find installation in default path
  117. def_path = Path(r'C:\Program Files\Git\bin\bash.exe')
  118. if def_path.is_file():
  119. return def_path
  120. # Find installation in PATH (C:\Program Files\Git\cmd should be in path)
  121. pth = (os.environ.get('PATH') or '').lower().split(';')
  122. pth = [p for p in pth if p.endswith(r'\git\cmd')]
  123. if pth:
  124. return Path(pth[0]).parent / r'bin\bash.exe'
  125. # Previously downloaded portable installation
  126. path = Path(__file__).parent / 'min_git'
  127. pkg_path = path / 'package.zip'
  128. if path.is_dir():
  129. return path / r'bin\bash.exe'
  130. # No installation found, download a portable installation
  131. print('Git installation not found. Git is required to use HyFetch/neofetch on Windows')
  132. print('Downloading a minimal portable package for Git...')
  133. urlretrieve(MINGIT_URL, pkg_path)
  134. print('Download finished! Extracting...')
  135. with zipfile.ZipFile(pkg_path, 'r') as zip_ref:
  136. zip_ref.extractall(path)
  137. print('Done!')
  138. return path / r'bin\bash.exe'
  139. def check_windows_cmd():
  140. """
  141. Check if this script is running under cmd.exe. If so, launch an external window with git bash
  142. since cmd doesn't support RGB colors.
  143. """
  144. if platform.system() == 'Windows':
  145. import psutil
  146. if psutil.Process(os.getppid()).name().lower().strip() == 'cmd.exe':
  147. print("cmd.exe doesn't support RGB colors, restarting in MinTTY...")
  148. cmd = f'"{ensure_git_bash().parent.parent / "usr/bin/mintty.exe"}" -s 110,40 -e python -m hyfetch --ask-exit'
  149. os.system(cmd)
  150. exit()
  151. def run_command(args: str, pipe: bool = False) -> str | None:
  152. """
  153. Run neofetch command
  154. """
  155. if platform.system() != 'Windows':
  156. full_cmd = shlex.split(f'{get_command_path()} {args}')
  157. else:
  158. cmd = get_command_path().replace("\\", "/").replace("C:/", "/c/")
  159. args = args.replace('\\', '/').replace('C:/', '/c/')
  160. full_cmd = [ensure_git_bash(), '-c', f'{cmd} {args}']
  161. # print(full_cmd)
  162. if pipe:
  163. return check_output(full_cmd).decode().strip()
  164. else:
  165. subprocess.run(full_cmd)
  166. def get_distro_ascii(distro: str | None = None) -> str:
  167. """
  168. Get the distro ascii of the current distro. Or if distro is specified, get the specific distro's
  169. ascii art instead.
  170. :return: Distro ascii
  171. """
  172. if not distro and GLOBAL_CFG.override_distro:
  173. distro = GLOBAL_CFG.override_distro
  174. if GLOBAL_CFG.debug:
  175. print(distro)
  176. print(GLOBAL_CFG)
  177. cmd = 'print_ascii'
  178. if distro:
  179. os.environ['CUSTOM_DISTRO'] = distro
  180. cmd = 'print_custom_ascii'
  181. return normalize_ascii(run_command(cmd, True))
  182. def get_distro_name():
  183. return run_command('ascii_distro_name', True)
  184. def run_neofetch(preset: ColorProfile, alignment: ColorAlignment):
  185. """
  186. Run neofetch with colors
  187. :param preset: Color palette
  188. :param alignment: Color alignment settings
  189. """
  190. asc = get_distro_ascii()
  191. w, h = ascii_size(asc)
  192. asc = alignment.recolor_ascii(asc, preset)
  193. # Write temp file
  194. with TemporaryDirectory() as tmp_dir:
  195. tmp_dir = Path(tmp_dir)
  196. path = tmp_dir / 'ascii.txt'
  197. path.write_text(asc)
  198. # Call neofetch with the temp file
  199. os.environ['ascii_len'] = str(w)
  200. os.environ['ascii_lines'] = str(h)
  201. run_command(f'--ascii --source {path.absolute()} --ascii-colors')
  202. def get_fore_back(distro: str | None = None) -> tuple[int, int] | None:
  203. """
  204. Get recommended foreground-background configuration for distro, or None if the distro ascii is
  205. not suitable for fore-back configuration.
  206. :return:
  207. """
  208. if not distro and GLOBAL_CFG.override_distro:
  209. distro = GLOBAL_CFG.override_distro
  210. if not distro:
  211. distro = get_distro_name().lower()
  212. for k, v in fore_back.items():
  213. if distro.startswith(k.lower()):
  214. return v
  215. return None
  216. # Foreground-background recommendation
  217. fore_back = {
  218. 'fedora': (2, 1),
  219. 'ubuntu': (2, 1),
  220. }