354 lines
12 KiB
Python
354 lines
12 KiB
Python
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||
#
|
||
# This program is free software: you can redistribute it and/or modify
|
||
# it under the terms of the GNU Affero General Public License as published
|
||
# by the Free Software Foundation, either version 3 of the License, or
|
||
# (at your option) any later version.
|
||
#
|
||
# This program is distributed in the hope that it will be useful,
|
||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||
# GNU Affero General Public License for more details.
|
||
#
|
||
# You should have received a copy of the GNU Affero General Public License
|
||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||
import logging
|
||
from typing import Any, Callable, List, Optional, Union
|
||
|
||
import pandas as pd
|
||
|
||
from utils.utils import colorText, get_sanitized_input
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class Selector:
|
||
@staticmethod
|
||
def _get_sorted_items(
|
||
items: List[Any], label_func: Callable[[Any], str]
|
||
) -> List[Any]:
|
||
return sorted(items, key=lambda item: label_func(item).lower())
|
||
|
||
@staticmethod
|
||
def _display_choices(
|
||
items: List[Any],
|
||
label_func: Callable[[Any], str],
|
||
num_columns: int = 3,
|
||
header: str = "Available Choices:",
|
||
) -> None:
|
||
# Force single column if items are DataFrame rows
|
||
|
||
if items and isinstance(items[0], (pd.Series, dict)):
|
||
num_columns = 1
|
||
|
||
rows = (len(items) + num_columns - 1) // num_columns
|
||
print(f"\n{header}")
|
||
for row in range(rows):
|
||
line = ""
|
||
for col in range(num_columns):
|
||
idx = row + col * rows
|
||
if idx < len(items):
|
||
label = label_func(items[idx])
|
||
line += f"{idx + 1}: {label:<30}"
|
||
print(line)
|
||
|
||
@staticmethod
|
||
def _display_selected_items(
|
||
selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3
|
||
) -> None:
|
||
print(colorText("\nCurrent selections:", "cyan"))
|
||
if not selected:
|
||
print(" (none)")
|
||
return
|
||
sorted_selected = sorted(selected, key=lambda item: label_func(item).lower())
|
||
rows = (len(sorted_selected) + num_columns - 1) // num_columns
|
||
for row in range(rows):
|
||
line = ""
|
||
for col in range(num_columns):
|
||
idx = row + col * rows
|
||
if idx < len(sorted_selected):
|
||
label = label_func(sorted_selected[idx])
|
||
line += f"{label:<30}"
|
||
print(line)
|
||
|
||
@staticmethod
|
||
def _parse_selection_input(input_str: str, max_index: int) -> List[int]:
|
||
selections = []
|
||
for part in input_str.split(","):
|
||
part = part.strip()
|
||
if "-" in part:
|
||
try:
|
||
start, end = map(int, part.split("-"))
|
||
selections.extend(range(start, end + 1))
|
||
except ValueError:
|
||
continue
|
||
elif part.isdigit():
|
||
selections.append(int(part))
|
||
return [i for i in selections if 1 <= i <= max_index]
|
||
|
||
@staticmethod
|
||
def _select_from_list(
|
||
items: List[Any],
|
||
label_func: Callable[[Any], str],
|
||
allow_multiple: bool = False,
|
||
prompt_each: bool = False,
|
||
header: str = "Available Choices:",
|
||
num_columns: int = 3,
|
||
) -> Union[Optional[Any], List[Any]]:
|
||
if not items:
|
||
logger.warning("No items available for selection.")
|
||
return None
|
||
|
||
full_sorted_items = Selector._get_sorted_items(items, label_func)
|
||
remaining_items = full_sorted_items.copy()
|
||
selected = []
|
||
|
||
if allow_multiple:
|
||
while True:
|
||
Selector._display_choices(
|
||
remaining_items, label_func, num_columns=num_columns, header=header
|
||
)
|
||
Selector._display_selected_items(
|
||
selected, label_func, num_columns=num_columns
|
||
)
|
||
choice = (
|
||
get_sanitized_input(
|
||
"Select item(s) by number (e.g. 1,3-5), R to reset, Q to finish: "
|
||
)
|
||
.strip()
|
||
.lower()
|
||
)
|
||
if choice == "q":
|
||
break
|
||
elif choice == "r":
|
||
selected.clear()
|
||
remaining_items = full_sorted_items.copy()
|
||
print(colorText("🔄 Selections reset.", "yellow"))
|
||
continue
|
||
indices = Selector._parse_selection_input(choice, len(remaining_items))
|
||
newly_selected = []
|
||
for index in indices:
|
||
item = remaining_items[index - 1]
|
||
if item not in selected:
|
||
selected.append(item)
|
||
newly_selected.append(item)
|
||
if prompt_each:
|
||
logger.info(f"Selected: {label_func(item)}")
|
||
else:
|
||
logger.warning("Item already selected.")
|
||
remaining_items = [
|
||
item for item in remaining_items if item not in newly_selected
|
||
]
|
||
return selected if selected else None
|
||
else:
|
||
Selector._display_choices(
|
||
full_sorted_items, label_func, num_columns=num_columns, header=header
|
||
)
|
||
try:
|
||
choice = int(get_sanitized_input("Select one item by number: "))
|
||
if 1 <= choice <= len(full_sorted_items):
|
||
selected_item = full_sorted_items[choice - 1]
|
||
logger.info(f"Selected: {label_func(selected_item)}")
|
||
return selected_item
|
||
else:
|
||
logger.warning("Selection out of range.")
|
||
except ValueError:
|
||
logger.warning("Invalid input.")
|
||
return None
|
||
|
||
@staticmethod
|
||
def select_with_mode(
|
||
items: List[Any],
|
||
label_func: Callable[[Any], str],
|
||
header: str = "Available Choices:",
|
||
) -> List[Any]:
|
||
print(
|
||
colorText(
|
||
"Choose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
|
||
"white",
|
||
)
|
||
)
|
||
mode = get_sanitized_input("").strip().lower()
|
||
if mode == "a":
|
||
return items
|
||
selected = Selector._select_from_list(
|
||
items,
|
||
label_func=label_func,
|
||
allow_multiple=True,
|
||
prompt_each=False,
|
||
header=header,
|
||
)
|
||
if not selected:
|
||
return items
|
||
if mode == "i":
|
||
print(colorText(f"✅ Included {len(selected)} item(s).", "green"))
|
||
return selected
|
||
elif mode == "e":
|
||
print(colorText(f"🚫 Excluded {len(selected)} item(s).", "yellow"))
|
||
return [item for item in items if item not in selected]
|
||
else:
|
||
print(colorText("⚠️ Invalid mode. Returning all items.", "yellow"))
|
||
return items
|
||
|
||
@staticmethod
|
||
def select_objects(
|
||
objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False
|
||
) -> Union[Optional[Any], List[Any]]:
|
||
return Selector._select_from_list(
|
||
objects,
|
||
label_func=lambda obj: getattr(obj, "name", str(obj)),
|
||
allow_multiple=allow_multiple,
|
||
prompt_each=prompt_each,
|
||
header="Available Objects:",
|
||
)
|
||
|
||
@staticmethod
|
||
def select_string(
|
||
options: List[str], allow_multiple: bool = False, prompt_each: bool = False
|
||
) -> Union[Optional[str], List[str]]:
|
||
return Selector._select_from_list(
|
||
options,
|
||
label_func=str,
|
||
allow_multiple=allow_multiple,
|
||
prompt_each=prompt_each,
|
||
header="Available Options:",
|
||
)
|
||
|
||
@staticmethod
|
||
def select_int(
|
||
options: List[int], allow_multiple: bool = False, prompt_each: bool = False
|
||
) -> Union[Optional[int], List[int]]:
|
||
return Selector._select_from_list(
|
||
options,
|
||
label_func=lambda x: str(x),
|
||
allow_multiple=allow_multiple,
|
||
prompt_each=prompt_each,
|
||
header="Available Integers:",
|
||
)
|
||
|
||
@staticmethod
|
||
def select_value(
|
||
prompt: str,
|
||
value_type: type = int,
|
||
valid_range: Optional[tuple] = None,
|
||
allow_quit: bool = False,
|
||
) -> Optional[Any]:
|
||
while True:
|
||
user_input = get_sanitized_input(prompt).strip().lower()
|
||
if allow_quit and user_input == "q":
|
||
logger.info("User opted to quit value selection.")
|
||
return None
|
||
try:
|
||
value = value_type(user_input)
|
||
if valid_range:
|
||
min_val, max_val = valid_range
|
||
if not (min_val <= value <= max_val):
|
||
logger.warning(f"Value out of range ({min_val}–{max_val}).")
|
||
continue
|
||
logger.info(f"User selected value: {value}")
|
||
return value
|
||
except ValueError:
|
||
logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
|
||
|
||
@staticmethod
|
||
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
|
||
while True:
|
||
response = get_sanitized_input(prompt).strip().lower()
|
||
if response in ["y", "yes"]:
|
||
logger.info("User confirmed action.")
|
||
return True
|
||
elif response in ["n", "no"]:
|
||
logger.info("User declined action.")
|
||
return False
|
||
else:
|
||
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
|
||
|
||
@staticmethod
|
||
def select_dataframe_rows(
|
||
df: pd.DataFrame,
|
||
columns: Optional[List[str]] = None,
|
||
allow_multiple: bool = False,
|
||
prompt_each: bool = False,
|
||
header: str = "Available Rows:",
|
||
) -> List[pd.Series]:
|
||
if df.empty:
|
||
print("DataFrame is empty.")
|
||
return []
|
||
|
||
if columns:
|
||
df = df[columns]
|
||
|
||
items = [row for _, row in df.iterrows()]
|
||
label_func = lambda row: str(row.to_dict())
|
||
|
||
result = Selector._select_from_list(
|
||
items,
|
||
label_func=label_func,
|
||
allow_multiple=allow_multiple,
|
||
prompt_each=prompt_each,
|
||
header=header,
|
||
)
|
||
|
||
if isinstance(result, pd.Series):
|
||
return [result]
|
||
elif isinstance(result, list):
|
||
return result
|
||
else:
|
||
return []
|
||
|
||
@staticmethod
|
||
def select_dataframe_with_mode(
|
||
df: pd.DataFrame,
|
||
columns: Optional[List[str]] = None,
|
||
header: str = "Available Rows:",
|
||
) -> List[pd.Series]:
|
||
if df.empty:
|
||
print("⚠️ DataFrame is empty.")
|
||
return []
|
||
|
||
# Filter columns if specified
|
||
if columns:
|
||
df = df[columns]
|
||
|
||
items = df.to_dict("records")
|
||
label_func = lambda row: " | ".join(str(row[col]) for col in df.columns)
|
||
|
||
# Show rows first
|
||
print(colorText(header, "cyan"))
|
||
for i, row in enumerate(items):
|
||
print(f"{i}: {label_func(row)}")
|
||
|
||
# Prompt for mode once
|
||
print(
|
||
colorText(
|
||
"\nChoose selection mode: [I]nclude only selected, [E]xclude selected, [A]ll (skip):",
|
||
"white",
|
||
)
|
||
)
|
||
mode = get_sanitized_input("").strip().lower()
|
||
|
||
if mode == "a":
|
||
return [pd.Series(row) for row in items]
|
||
|
||
# Prompt for selection only once
|
||
selected = Selector._select_from_list(
|
||
items,
|
||
label_func=label_func,
|
||
allow_multiple=True,
|
||
prompt_each=False,
|
||
header=header,
|
||
)
|
||
|
||
if not selected:
|
||
return [pd.Series(row) for row in items]
|
||
|
||
if mode == "i":
|
||
print(colorText(f"✅ Included {len(selected)} row(s).", "green"))
|
||
return [pd.Series(row) for row in selected]
|
||
elif mode == "e":
|
||
print(colorText(f"🚫 Excluded {len(selected)} row(s).", "yellow"))
|
||
return [pd.Series(row) for row in items if row not in selected]
|
||
else:
|
||
print(colorText("⚠️ Invalid mode. Returning no rows.", "yellow"))
|
||
return []
|