neofetch_util.py 6.1 KB

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