refactor: replace VfbConnect with SolrTermInfoFetcher and enhance DataFrame support

Author: Robbie1977

src/vfbquery/solr_fetcher.py

@@ -1,7 +1,9 @@
 import requests
 import json
 import logging
-from typing import List, Dict, Any, Optional
+import pandas as pd
+from typing import List, Dict, Any, Optional, Union
+from vfb_connect import vfb
 
 class SolrTermInfoFetcher:
     """Fetches term information directly from the Solr server instead of using VfbConnect"""
@@ -10,20 +12,21 @@ def __init__(self, solr_url: str = "https://solr.virtualflybrain.org/solr/vfb_js
         """Initialize with the Solr server URL"""
         self.solr_url = solr_url
         self.logger = logging.getLogger(__name__)
+        self.vfb = vfb
 
     def get_TermInfo(self, short_forms: List[str], 
                     return_dataframe: bool = False, 
-                    summary: bool = False) -> List[Dict[str, Any]]:
+                    summary: bool = False) -> Union[List[Dict[str, Any]], pd.DataFrame]:
         """
         Fetch term info from Solr directly, mimicking VFBconnect's interface
         
         Args:
             short_forms: List of term IDs to fetch
-            return_dataframe: If True, return as pandas DataFrame (not fully implemented)
+            return_dataframe: If True, return as pandas DataFrame
             summary: If True, return summarized version
             
         Returns:
-            List of term info dictionaries
+            List of term info dictionaries or DataFrame
         """
         results = []
 
@@ -54,7 +57,6 @@ def get_TermInfo(self, short_forms: List[str],
 
                 # Extract and parse the term_info string which is itself JSON
                 term_info_str = docs[0]["term_info"][0]
-                # No need to handle escapes - json.loads does that automatically
                 term_info_obj = json.loads(term_info_str)
                 results.append(term_info_obj)
 
@@ -65,9 +67,23 @@ def get_TermInfo(self, short_forms: List[str],
             except Exception as e:
                 self.logger.error(f"Unexpected error for {short_form}: {e}")
 
-        # Handle dataframe conversion if needed (this would need to be implemented)
-        if return_dataframe:
-            self.logger.warning("return_dataframe=True not fully implemented")
-            # You would need to implement pandas DataFrame conversion logic here
+        # Convert to DataFrame if requested
+        if return_dataframe and results:
+            try:
+                return pd.json_normalize(results)
+            except Exception as e:
+                self.logger.error(f"Error converting to DataFrame: {e}")
+                return results
 
-        return results
+        return results
+    
+    # Pass through any non-implemented methods to VFBconnect
+    def __getattr__(self, name):
+        """
+        Automatically pass through any non-implemented methods to VFBconnect
+        
+        This allows us to use this class as a drop-in replacement for VfbConnect
+        while only implementing the methods we want to customize.
+        """
+        self.logger.debug(f"Passing through method call: {name}")
+        return getattr(self.vfb, name)

src/vfbquery/vfb_queries.py

@@ -1,6 +1,9 @@
 import pysolr
 from .term_info_queries import deserialize_term_info
-from vfb_connect.cross_server_tools import VfbConnect, dict_cursor
+# Replace VfbConnect import with our new SolrTermInfoFetcher
+from .solr_fetcher import SolrTermInfoFetcher
+# Keep dict_cursor if it's used elsewhere
+from vfb_connect.cross_server_tools import dict_cursor
 from marshmallow import Schema, fields, post_load
 from typing import List, Tuple, Dict, Any, Union
 import pandas as pd
@@ -10,8 +13,8 @@
 # Connect to the VFB SOLR server
 vfb_solr = pysolr.Solr('http://solr.virtualflybrain.org/solr/vfb_json/', always_commit=False, timeout=990)
 
-# Create a VFB connection object for retrieving instances
-vc = VfbConnect()
+# Replace VfbConnect with SolrTermInfoFetcher
+vc = SolrTermInfoFetcher()
 
 class Query:
     def __init__(self, query, label, function, takes, preview=0, preview_columns=[], preview_results=[], output_format="table", count=-1):