pen-tracker/build/lib/pen_tracker/engine.py
2026-05-05 10:41:23 -05:00

163 lines
6.4 KiB
Python

import csv
import os
import logging
from dataclasses import dataclass, asdict
from typing import List
logger = logging.getLogger(__name__)
@dataclass
class Ink:
Vendor: str = "N/A"
Name: str = "N/A"
Color: str = "N/A"
Purchased: str = "N/A"
Size: str = "N/A"
Notes: str = "N/A"
@dataclass
class Pen:
Make: str = "N/A"
Model: str = "N/A"
Date_Purchased: str = "N/A"
Vendor: str = "N/A"
Nib: str = "N/A"
Nib_Material: str = "N/A"
Body: str = "N/A"
Cap: str = "N/A"
Post: str = "N/A"
Current_Ink: str = "N/A"
Inked_date: str = "N/A"
Notes: str = "N/A"
class InkTracker:
def __init__(self, storage_file: str = None):
if storage_file is None:
storage_file = os.getenv('INK_TRACKER_CSV')
if storage_file is None:
data_home = os.getenv('XDG_DATA_HOME', os.path.expanduser('~/.local/share'))
app_data_dir = os.path.join(data_home, 'pen-tracker')
os.makedirs(app_data_dir, exist_ok=True)
storage_file = os.path.join(app_data_dir, 'inks.csv')
self.storage_file = storage_file
self.headers = ['Vendor', 'Name', 'Color', 'Purchased', 'Size', 'Notes']
self.inks: List[Ink] = self.load_data()
def _sort_inks(self):
"""Sorts the inks list by Vendor, then by Name alphabetically."""
self.inks.sort(key=lambda x: (x.Vendor.lower(), x.Name.lower()))
def load_data(self) -> List[Ink]:
"""Loads data from the CSV file."""
if not os.path.exists(self.storage_file):
self._create_empty_csv()
return []
inks = []
try:
with open(self.storage_file, mode='r', encoding='utf-8-sig') as f:
reader = csv.DictReader(f)
if reader.fieldnames:
reader.fieldnames = [h.strip() for h in reader.fieldnames]
for row in reader:
if None in row:
extras = row.pop(None)
if extras:
extra_text = ",".join(extras).strip()
if extra_text:
row['Notes'] = (row.get('Notes') or '')
if row['Notes']:
row['Notes'] += ", "
row['Notes'] += extra_text
clean_row = {
k.strip(): (v.strip() if v else "N/A")
for k, v in row.items() if k is not None
}
inks.append(Ink(**clean_row))
inks.sort(key=lambda x: (x.Vendor.lower(), x.Name.lower()))
except Exception as e:
logger.error(f"Error loading inks CSV: {e}")
return inks
def _create_empty_csv(self):
"""Creates the CSV file with headers if it is missing."""
with open(self.storage_file, mode='w', newline='', encoding='utf-8') as f:
writer = csv.DictWriter(f, fieldnames=self.headers)
writer.writeheader()
def save_data(self):
"""Sorts the data before writing to ensure persistent alphabetical order."""
self._sort_inks()
with open(self.storage_file, mode='w', newline='', encoding='utf-8') as f:
writer = csv.DictWriter(f, fieldnames=self.headers)
writer.writeheader()
for ink in self.inks:
writer.writerow(asdict(ink))
class PenTracker:
def __init__(self, storage_file: str = None):
if storage_file is None:
storage_file = os.getenv('PEN_TRACKER_CSV')
if storage_file is None:
data_home = os.getenv('XDG_DATA_HOME', os.path.expanduser('~/.local/share'))
app_data_dir = os.path.join(data_home, 'pen-tracker')
os.makedirs(app_data_dir, exist_ok=True)
storage_file = os.path.join(app_data_dir, 'pens.csv')
self.storage_file = storage_file
self.headers = [
'Make', 'Model', 'Date-Purchased', 'Vendor', 'Nib',
'Nib-Material', 'Body', 'Cap', 'Post',
'Current-Ink', 'Inked-date', 'Notes'
]
self.key_map = {
'Date-Purchased': 'Date_Purchased',
'Inked-date': 'Inked_date',
'Nib-Material': 'Nib_Material',
'Current-Ink': 'Current_Ink'
}
self.reverse_key_map = {v: k for k, v in self.key_map.items()}
self.pens: List[Pen] = self.load_data()
def _sort_pens(self):
"""Sorts the pens list by Make, then by Model alphabetically."""
self.pens.sort(key=lambda x: (x.Make.lower(), x.Model.lower()))
def load_data(self) -> List[Pen]:
"""Loads data from the CSV file."""
if not os.path.exists(self.storage_file):
self._create_empty_csv()
return []
pens = []
try:
# Standard, clean way to open the file
with open(self.storage_file, mode='r', encoding='utf-8-sig') as f:
reader = csv.DictReader(f)
for row in reader:
# Strip whitespace from keys and values, handle empty values
clean_row = {k.strip(): (v.strip() if v else "N/A") for k, v in row.items()}
# Map keys to dataclass field names
mapped_row = {self.key_map.get(k, k): v for k, v in clean_row.items()}
pens.append(Pen(**mapped_row))
pens.sort(key=lambda x: (x.Make.lower(), x.Model.lower()))
except Exception as e:
logger.error(f"Error loading CSV: {e}")
return pens
def _create_empty_csv(self):
"""Creates the CSV file with headers if it is missing."""
with open(self.storage_file, mode='w', newline='', encoding='utf-8') as f:
writer = csv.DictWriter(f, fieldnames=self.headers)
writer.writeheader()
def save_data(self):
"""Sorts the data before writing to ensure persistent alphabetical order."""
self._sort_pens()
with open(self.storage_file, mode='w', newline='', encoding='utf-8') as f:
writer = csv.DictWriter(f, fieldnames=self.headers)
writer.writeheader()
for pen in self.pens:
row = asdict(pen)
# Map back to CSV keys
csv_row = {k.replace('_', '-'): v for k, v in row.items()}
writer.writerow(csv_row)