neofetch_util.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. from __future__ import annotations
  2. import inspect
  3. import os
  4. import platform
  5. import re
  6. import subprocess
  7. from dataclasses import dataclass
  8. from pathlib import Path
  9. from subprocess import check_output
  10. from tempfile import TemporaryDirectory
  11. import pkg_resources
  12. from typing_extensions import Literal
  13. from hyfetch.color_util import color
  14. from .constants import GLOBAL_CFG
  15. from .presets import ColorProfile
  16. from .serializer import from_dict
  17. RE_NEOFETCH_COLOR = re.compile('\\${c[0-9]}')
  18. def term_size() -> tuple[int, int]:
  19. """
  20. Get terminal size
  21. :return:
  22. """
  23. try:
  24. return os.get_terminal_size().columns, os.get_terminal_size().lines
  25. except Exception:
  26. return 100, 20
  27. def ascii_size(asc: str) -> tuple[int, int]:
  28. """
  29. Get distro ascii width, height ignoring color code
  30. :param asc: Distro ascii
  31. :return: Width, Height
  32. """
  33. return max(len(line) for line in re.sub(RE_NEOFETCH_COLOR, '', asc).split('\n')), len(asc.split('\n'))
  34. def normalize_ascii(asc: str) -> str:
  35. """
  36. Make sure every line are the same width
  37. """
  38. w = ascii_size(asc)[0]
  39. return '\n'.join(line + ' ' * (w - ascii_size(line)[0]) for line in asc.split('\n'))
  40. def fill_starting(asc: str) -> str:
  41. """
  42. Fill the missing starting placeholders.
  43. E.g. "${c1}...\n..." -> "${c1}...\n${c1}..."
  44. """
  45. new = []
  46. last = ''
  47. for line in asc.split('\n'):
  48. new.append(last + line)
  49. # Line has color placeholders
  50. matches = RE_NEOFETCH_COLOR.findall(line)
  51. if len(matches) > 0:
  52. # Get the last placeholder for the next line
  53. last = matches[-1]
  54. return '\n'.join(new)
  55. @dataclass
  56. class ColorAlignment:
  57. mode: Literal['horizontal', 'vertical', 'custom']
  58. # custom_colors[ascii color index] = unique color index in preset
  59. custom_colors: dict[int, int] = ()
  60. # Foreground/background ascii color index
  61. fore_back: tuple[int, int] = ()
  62. @classmethod
  63. def from_dict(cls, d: dict):
  64. return from_dict(cls, d)
  65. def recolor_ascii(self, asc: str, preset: ColorProfile) -> str:
  66. """
  67. Use the color alignment to recolor an ascii art
  68. :return Colored ascii, Uncolored lines
  69. """
  70. asc = fill_starting(asc)
  71. if self.fore_back and self.mode in ['horizontal', 'vertical']:
  72. fore, back = self.fore_back
  73. # Replace foreground colors
  74. asc = asc.replace(f'${{c{fore}}}', color('&0' if GLOBAL_CFG.is_light else '&f'))
  75. lines = asc.split('\n')
  76. # Add new colors
  77. if self.mode == 'horizontal':
  78. colors = preset.with_length(len(lines))
  79. asc = '\n'.join([l.replace(f'${{c{back}}}', colors[i].to_ansi()) + color('&r') for i, l in enumerate(lines)])
  80. else:
  81. raise NotImplementedError()
  82. # Remove existing colors
  83. asc = re.sub(RE_NEOFETCH_COLOR, '', asc)
  84. elif self.mode in ['horizontal', 'vertical']:
  85. # Remove existing colors
  86. asc = re.sub(RE_NEOFETCH_COLOR, '', asc)
  87. lines = asc.split('\n')
  88. # Add new colors
  89. if self.mode == 'horizontal':
  90. colors = preset.with_length(len(lines))
  91. asc = '\n'.join([colors[i].to_ansi() + l + color('&r') for i, l in enumerate(lines)])
  92. else:
  93. asc = '\n'.join(preset.color_text(line) + color('&r') for line in lines)
  94. else:
  95. preset = preset.unique_colors()
  96. # Apply colors
  97. color_map = {ai: preset.colors[pi].to_ansi() for ai, pi in self.custom_colors.items()}
  98. for ascii_i, c in color_map.items():
  99. asc = asc.replace(f'${{c{ascii_i}}}', c)
  100. return asc
  101. def get_command_path() -> str:
  102. """
  103. Get the absolute path of the neofetch command
  104. :return: Command path
  105. """
  106. return pkg_resources.resource_filename(__name__, 'scripts/neowofetch')
  107. def get_distro_ascii(distro: str | None = None) -> str:
  108. """
  109. Get the distro ascii of the current distro. Or if distro is specified, get the specific distro's
  110. ascii art instead.
  111. :return: Distro ascii
  112. """
  113. if not distro and GLOBAL_CFG.override_distro:
  114. distro = GLOBAL_CFG.override_distro
  115. if GLOBAL_CFG.debug:
  116. print(distro)
  117. print(GLOBAL_CFG)
  118. cmd = 'print_ascii'
  119. if distro:
  120. os.environ['CUSTOM_DISTRO'] = distro
  121. cmd = 'print_custom_ascii'
  122. return normalize_ascii(check_output([get_command_path(), cmd]).decode().strip())
  123. def get_distro_name():
  124. return check_output([get_command_path(), 'ascii_distro_name']).decode().strip()
  125. def run_neofetch(preset: ColorProfile, alignment: ColorAlignment):
  126. """
  127. Run neofetch with colors
  128. :param preset: Color palette
  129. :param alignment: Color alignment settings
  130. """
  131. asc = get_distro_ascii()
  132. w, h = ascii_size(asc)
  133. asc = alignment.recolor_ascii(asc, preset)
  134. # Write temp file
  135. with TemporaryDirectory() as tmp_dir:
  136. tmp_dir = Path(tmp_dir)
  137. path = tmp_dir / 'ascii.txt'
  138. path.write_text(asc)
  139. # Call neofetch with the temp file
  140. os.environ['ascii_len'] = str(w)
  141. os.environ['ascii_lines'] = str(h)
  142. if platform.system() != 'Windows':
  143. os.system(f'{get_command_path()} --ascii --source {path.absolute()} --ascii-colors')
  144. else:
  145. cmd = get_command_path().replace("\\", "/").replace("C:/", "/c/")
  146. path_str = str(path.absolute()).replace('\\', '/').replace('C:/', '/c/')
  147. cmd = f'ascii_len={w} ascii_lines={h} {cmd} --ascii --source {path_str} --ascii-colors'
  148. full_cmd = ['C:\\Program Files\\Git\\bin\\bash.exe', '-c', cmd]
  149. # print(full_cmd)
  150. subprocess.run(full_cmd)
  151. def get_fore_back(distro: str | None = None) -> tuple[int, int] | None:
  152. """
  153. Get recommended foreground-background configuration for distro, or None if the distro ascii is
  154. not suitable for fore-back configuration.
  155. :return:
  156. """
  157. if not distro and GLOBAL_CFG.override_distro:
  158. distro = GLOBAL_CFG.override_distro
  159. if not distro:
  160. distro = get_distro_name().lower()
  161. for k, v in fore_back.items():
  162. if distro.startswith(k.lower()):
  163. return v
  164. return None
  165. # Foreground-background recommendation
  166. fore_back = {
  167. 'fedora': (2, 1),
  168. 'ubuntu': (2, 1),
  169. }