neofetch_util.py 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. from __future__ import annotations
  2. import os
  3. import platform
  4. import re
  5. import subprocess
  6. from dataclasses import dataclass
  7. from pathlib import Path
  8. from subprocess import check_output
  9. from tempfile import TemporaryDirectory
  10. import pkg_resources
  11. from hyfetch.color_util import color
  12. from typing_extensions import Literal
  13. from .presets import ColorProfile
  14. RE_NEOFETCH_COLOR = re.compile('\\${c[0-9]}')
  15. def ascii_size(asc: str) -> tuple[int, int]:
  16. """
  17. Get distro ascii width, height ignoring color code
  18. :param asc: Distro ascii
  19. :return: Width, Height
  20. """
  21. return max(len(line) for line in re.sub(RE_NEOFETCH_COLOR, '', asc).split('\n')), len(asc.split('\n'))
  22. def normalize_ascii(asc: str) -> str:
  23. """
  24. Make sure every line are the same width
  25. """
  26. w = ascii_size(asc)[0]
  27. return '\n'.join(line + ' ' * (w - ascii_size(line)[0]) for line in asc.split('\n'))
  28. @dataclass
  29. class ColorAlignment:
  30. mode: Literal['horizontal', 'vertical', 'custom']
  31. # custom_colors[ascii color index] = unique color index in preset
  32. custom_colors: dict[int, int] = ()
  33. def recolor_ascii(self, asc: str, preset: ColorProfile) -> str:
  34. """
  35. Use the color alignment to recolor an ascii art
  36. :return Colored ascii, Uncolored lines
  37. """
  38. if self.mode in ['horizontal', 'vertical']:
  39. # Remove existing colors
  40. asc = re.sub(RE_NEOFETCH_COLOR, '', asc)
  41. lines = asc.split('\n')
  42. # Add new colors
  43. if self.mode == 'horizontal':
  44. colors = preset.with_length(len(lines))
  45. asc = '\n'.join([colors[i].to_ansi() + l + color('&r') for i, l in enumerate(lines)])
  46. else:
  47. asc = '\n'.join(preset.color_text(line) + color('&r') for line in lines)
  48. else:
  49. preset = preset.unique_colors()
  50. # Apply colors
  51. new = []
  52. start_color = None
  53. color_map = {ai: preset.colors[pi].to_ansi() for ai, pi in self.custom_colors.items()}
  54. for line in asc.split('\n'):
  55. # Line has color placeholders
  56. if len(RE_NEOFETCH_COLOR.findall(line)) > 0:
  57. # Get the last placeholder for the next line
  58. last = int(RE_NEOFETCH_COLOR.findall(line)[-1][3])
  59. # Replace placeholders
  60. for ascii_i, c in color_map.items():
  61. line = line.replace(f'${{c{ascii_i}}}', c)
  62. # Add to new ascii
  63. new.append(f'{start_color or ""}{line}')
  64. # Change next start color
  65. start_color = color_map[last]
  66. else:
  67. new.append(f'{start_color or ""}{line}')
  68. asc = '\n'.join(new)
  69. return asc
  70. def get_command_path() -> str:
  71. """
  72. Get the absolute path of the neofetch command
  73. :return: Command path
  74. """
  75. return pkg_resources.resource_filename(__name__, 'scripts/neofetch_mod.sh')
  76. def get_distro_ascii(distro: str | None = None) -> str:
  77. """
  78. Get the distro ascii of the current distro. Or if distro is specified, get the specific distro's
  79. ascii art instead.
  80. :return: Distro ascii
  81. """
  82. cmd = 'print_ascii'
  83. if distro:
  84. os.environ['CUSTOM_DISTRO'] = distro
  85. cmd = 'print_custom_ascii'
  86. return normalize_ascii(check_output([get_command_path(), cmd]).decode().strip())
  87. def run_neofetch(preset: ColorProfile, alignment: ColorAlignment):
  88. asc = get_distro_ascii()
  89. w, h = ascii_size(asc)
  90. asc = alignment.recolor_ascii(asc, preset)
  91. # Write temp file
  92. with TemporaryDirectory() as tmp_dir:
  93. tmp_dir = Path(tmp_dir)
  94. path = tmp_dir / 'ascii.txt'
  95. path.write_text(asc)
  96. # Call neofetch with the temp file
  97. os.environ['ascii_len'] = str(w)
  98. os.environ['ascii_lines'] = str(h)
  99. if platform.system() != 'Windows':
  100. os.system(f'{get_command_path()} --ascii --source {path.absolute()} --ascii-colors')
  101. if platform.system() == 'Windows':
  102. cmd = get_command_path().replace("\\", "/").replace("C:/", "/c/")
  103. path_str = str(path.absolute()).replace('\\', '/').replace('C:/', '/c/')
  104. cmd = f'ascii_len={w} ascii_lines={h} {cmd} --ascii --source {path_str} --ascii-colors'
  105. full_cmd = ['C:\\Program Files\\Git\\bin\\bash.exe', '-c', cmd]
  106. # print(full_cmd)
  107. subprocess.run(full_cmd)