kiwi-scp/kiwi_scp/wstring.py
2021-11-27 18:25:38 +01:00

97 lines
2.5 KiB
Python

import re
from enum import Enum, auto
from typing import List
import attr
import wcwidth
class WAlignment(Enum):
LEFT = auto()
RIGHT = auto()
CENTER = auto()
@attr.s
class WString:
s: str = attr.ib()
# from https://stackoverflow.com/a/38662876
ANSI_ESCAPES = re.compile(r"(?:\x1B[@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
def __str__(self) -> str:
return self.s
def __len__(self) -> int:
return wcwidth.wcswidth(WString.ANSI_ESCAPES.sub("", self.s))
def pad(self, alignment: WAlignment = WAlignment.CENTER, wlen: int = 0, char: str = " ") -> "WString":
char = char[0]
if alignment is WAlignment.LEFT:
return WString(f"{self}{char * wlen}")
elif alignment is WAlignment.RIGHT:
return WString(f"{char * wlen}{self}")
else:
pad_l, pad_r = wlen // 2, wlen - (wlen // 2)
return WString(f"{char * pad_l}{self}{char * pad_r}")
@attr.s
class WParagraph:
lines: List[WString] = attr.ib()
def __str__(self) -> str:
return "\n".join(
str(line)
for line in self.lines
)
@classmethod
def from_strings(cls, *source: str) -> "WParagraph":
return cls([
WString(line)
for line in source
])
def align(self, alignment: WAlignment = WAlignment.CENTER, padding: int = 0, char: str = " ") -> "WParagraph":
total_length = max(
len(line)
for line in self.lines
) + padding
pad_lengths = (
total_length - len(line)
for line in self.lines
)
return WParagraph([
line.pad(alignment, wlen, char)
for line, wlen in zip(self.lines, pad_lengths)
])
def surround(self, char: str, padding: int = 1) -> "WParagraph":
char = char[0]
padding = " " * padding
l_border, r_border = char + padding, padding + char
lines = [
WString(f"{l_border}{line}{r_border}")
for line in self.lines
]
extra_line = char * len(lines[0])
return WParagraph([
extra_line,
*lines,
extra_line,
])
def emphasize(self, count: int = 3, padding: int = 1) -> "WParagraph":
padding = " " * padding
l_border, r_border = (">" * count) + padding, padding + ("<" * count)
return WParagraph([
WString(f"{l_border}{line}{r_border}")
for line in self.lines
])