Update server.py

Author: jeffreyaven

pystackql/magic_ext/server.py

@@ -69,8 +69,12 @@ def _display_with_csv_download(self, df):
         
         :param df: The DataFrame to display and make downloadable.
         """
+        import IPython.display
+        
         try:
             # Generate CSV data
+            import io
+            import base64
             csv_buffer = io.StringIO()
             df.to_csv(csv_buffer, index=False)
             csv_data = csv_buffer.getvalue()
@@ -82,7 +86,7 @@ def _display_with_csv_download(self, df):
             download_link = f'data:text/csv;base64,{csv_base64}'
 
             # Display the DataFrame first
-            display(df)
+            IPython.display.display(df)
 
             # Create and display the download button
             download_html = f'''
@@ -95,15 +99,15 @@ def _display_with_csv_download(self, df):
                 </a>
             </div>
             '''
-            display(HTML(download_html))
+            IPython.display.display(IPython.display.HTML(download_html))
 
         except Exception as e:
             # If CSV generation fails, just display the DataFrame normally
-            display(df)
+            IPython.display.display(df)
             print(f"Error generating CSV download: {e}")
-
+        
 def load_ipython_extension(ipython):
     """Load the server magic in IPython."""
     # Create an instance of the magic class and register it
     magic_instance = StackqlServerMagic(ipython)
-    ipython.register_magics(magic_instance)
+    ipython.register_magics(magic_instance)