First commit

baloo_tools/baloo_tools.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+"""
+Baloo Tools Library
+Helper functions to interact directly with the Baloo LMDB index.
+"""
+
+import json
+import lmdb
+import os
+import sys
+from typing import Tuple
+
+
+class BalooTools:
+    """Class to interact directly with the Baloo LMDB index."""
+
+    def __init__(self) -> None:
+        """Initializes the connection path to the Baloo index."""
+        self.baloo_db_path = os.path.join(
+            os.path.expanduser("~"), ".local/share/baloo/index"
+        )
+
+    def get_resolution(self, file_id: int, sep: str = 'x') -> Tuple[int, int]:
+        """
+        Retrieves the width and height of an image/video from the Baloo index.
+
+        Args:
+            file_id: The integer ID of the file.
+            sep: Separator used (unused currently, kept for compatibility).
+
+        Returns:
+            A tuple of (width, height) integers. Returns (-1, -1) if not found.
+        """
+        try:
+            # Using context manager ensures the environment is closed properly
+            with lmdb.Environment(
+                self.baloo_db_path,
+                subdir=False,
+                readonly=True,
+                lock=False,
+                max_dbs=20
+            ) as env:
+                document_data_db = env.open_db(b'documentdatadb')
+
+                with env.begin() as txn:
+                    cursor = txn.cursor(document_data_db)
+
+                    # Convert ID to 8-byte little-endian format
+                    file_id_bytes = int.to_bytes(
+                        file_id, length=8, byteorder='little', signed=False
+                    )
+
+                    if cursor.set_range(file_id_bytes):
+                        for key, value in cursor:
+                            if key != file_id_bytes:
+                                break
+
+                            try:
+                                jvalue = json.loads(value.decode())
+                                # Baloo stores width in '26' and height in '27'
+                                return jvalue.get('26', -1), jvalue.get('27', -1)
+                            except (json.JSONDecodeError, KeyError):
+                                return -1, -1
+
+        except lmdb.Error as e:
+            print(f"Warning: Failed to access Baloo LMDB index: {e}", file=sys.stderr)
+
+        return -1, -1
+
+
+# Helper function to maintain compatibility with bagheera_search_lib.py
+# since it imports `get_resolution` directly.
+def get_resolution(file_id: int, sep: str = 'x') -> Tuple[int, int]:
+    """Standalone helper function to instantiate BalooTools and get resolution."""
+    tools = BalooTools()
+    return tools.get_resolution(file_id, sep)
+
+
+if __name__ == '__main__':
+    # CLI execution support for testing
+    if len(sys.argv) > 1:
+        try:
+            target_id = int(sys.argv[1], 16)
+            width, height = get_resolution(target_id)
+            print(f"{width} {height}")
+        except ValueError:
+            print("Error: Please provide a valid hexadecimal file ID.", file=sys.stderr)
+            sys.exit(1)