#!/usr/bin/env python3 import os import sqlite3 import subprocess import sys import tempfile # Configuration DB_PATH = os.path.expanduser("~/screenshot_ocr.db") SCREENSHOTS_DIR = os.path.expanduser("~/Screenshots") ROFI_PROMPT = "OCR Search" def check_dependencies(): """Check if required dependencies are installed.""" # Check if rofi is installed try: subprocess.run( ["which", "rofi"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) except subprocess.CalledProcessError: print( "Error: rofi is not installed. Please install it with: sudo pacman -S rofi" ) return False # Check if database exists if not os.path.exists(DB_PATH): print(f"Error: Database not found at {DB_PATH}") print("Run the OCR script first to create and populate the database.") return False return True def format_text(text, max_length=100): """Format OCR text for display in rofi.""" # Replace newlines with spaces text = text.replace("\n", " ") # Replace multiple spaces with a single space while " " in text: text = text.replace(" ", " ") # Truncate if too long if len(text) > max_length: text = text[:max_length] + "..." return text def get_screenshot_data(): """Get data from the SQLite database.""" try: conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Get all records cursor.execute( "SELECT filename, ocr_text, full_path FROM ocr_results ORDER BY filename" ) results = [] for filename, ocr_text, full_path in cursor.fetchall(): formatted_text = format_text(ocr_text) # Format for display: filename | ocr_text display_text = f"{filename} | {formatted_text}" results.append((display_text, full_path)) conn.close() return results except sqlite3.Error as e: print(f"Database error: {e}") return [] def show_rofi_menu(items): """Display rofi menu with the given items and return selection.""" if not items: print("No items to display.") return None, None # Create temporary file with menu items with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: for display_text, _ in items: temp_file.write(f"{display_text}\n") temp_file_path = temp_file.name try: # Run rofi command cmd = [ "rofi", "-dmenu", "-i", # Case-insensitive matching "-p", ROFI_PROMPT, "-width", "80", "-lines", "15", "-font", "mono 10", ] process = subprocess.Popen( cmd, stdin=open(temp_file_path, "r"), stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, ) stdout, stderr = process.communicate() if process.returncode != 0: print(f"Rofi error: {stderr}") return None, None selection = stdout.strip() if not selection: return None, None # Find the matching item for i, (display_text, full_path) in enumerate(items): if display_text == selection: return display_text, full_path return None, None finally: # Clean up temporary file os.unlink(temp_file_path) def open_file(file_path): """Open the selected file with the default application.""" if os.path.exists(file_path): try: subprocess.Popen( ["xdg-open", file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE ) print(f"Opening: {file_path}") return True except Exception as e: print(f"Error opening file: {e}") return False else: print(f"Error: File not found: {file_path}") return False def main(): """Main function.""" if not check_dependencies(): sys.exit(1) # Get screenshot data items = get_screenshot_data() if not items: print("No OCR data found in the database.") sys.exit(1) # Show rofi menu selection, file_path = show_rofi_menu(items) # Open selected file if selection and file_path: open_file(file_path) else: print("No selection made.") if __name__ == "__main__": main()