52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
from sawaw.sentiments import SentimentResult
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
import colorama
|
|
|
|
color_mapping = {
|
|
SentimentResult.POSITIVE: colorama.Fore.GREEN,
|
|
SentimentResult.NEUTRAL: colorama.Fore.YELLOW,
|
|
SentimentResult.NEGATIVE: colorama.Fore.RED,
|
|
SentimentResult.NONE: colorama.Fore.LIGHTWHITE_EX,
|
|
SentimentResult.UNDEFINED: colorama.Fore.MAGENTA,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class SAWAWEntry:
|
|
comment: str
|
|
aspect_words: List[str]
|
|
sentiment_results: List[SentimentResult] = None
|
|
|
|
def __post_init__(self):
|
|
if self.sentiment_results is None:
|
|
self.sentiment_results = [SentimentResult.UNDEFINED] * len(
|
|
self.aspect_words
|
|
)
|
|
|
|
def __repr__(self):
|
|
colored_comment = self.comment
|
|
for aspect_word in self.aspect_words:
|
|
colored_comment = colored_comment.replace(
|
|
aspect_word,
|
|
color_mapping[
|
|
self.sentiment_results[self.aspect_words.index(aspect_word)]
|
|
]
|
|
+ aspect_word
|
|
+ colorama.Style.RESET_ALL,
|
|
)
|
|
return colored_comment
|
|
|
|
def __str__(self):
|
|
return self.__repr__()
|
|
|
|
@staticmethod
|
|
def print_legends():
|
|
legends = []
|
|
for sentiment_result, color in color_mapping.items():
|
|
legends.append(
|
|
"{}{}{}".format(color, sentiment_result.name, colorama.Style.RESET_ALL)
|
|
)
|
|
print(" | ".join(legends))
|
|
|