Fix stack traces on exit during script (#176)

Author: brendan-priorlabs

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tabpfn-client"
-version = "0.2.7"
+version = "0.2.8"
 requires-python = ">=3.9"
 dynamic = ["dependencies", "optional-dependencies"]
 

tabpfn_client/prompt_agent.py

@@ -1,6 +1,7 @@
 #  Copyright (c) Prior Labs GmbH 2025.
 #  Licensed under the Apache License, Version 2.0
 import getpass
+import sys
 import textwrap
 from rich.table import Table
 
@@ -17,6 +18,17 @@
 )
 
 
+def maybe_graceful_exit() -> None:
+    try:
+        from IPython import get_ipython
+
+        if get_ipython() is not None:
+            return
+    except ImportError:
+        # We're in a script, just exit
+        sys.exit(1)
+
+
 class PromptAgent:
     def __new__(cls):
         raise RuntimeError(
@@ -117,6 +129,7 @@ def prompt_and_set_token(cls) -> bool:
             return True
         except KeyboardInterrupt:
             console.print("\n\n[yellow]Interrupted. Goodbye![/yellow]")
+            maybe_graceful_exit()
             return False
 
     @classmethod
@@ -145,6 +158,7 @@ def _prompt_and_set_token_impl(cls) -> bool:
 
         if choice == "q":
             console.print("Goodbye!")
+            maybe_graceful_exit()
             return False
 
         # Registration
@@ -347,6 +361,7 @@ def _prompt_and_set_token_impl(cls) -> bool:
                     continue
                 elif retry_choice == "q":
                     console.print("Goodbye!")
+                    maybe_graceful_exit()
                     return False
                 else:
                     # Invalid choice, use default (retry)
@@ -551,7 +566,8 @@ def reverify_email(cls, access_token):
                 return "restart"  # Signal to show main menu
             elif choice in ["q", "quit"]:
                 console.print("Goodbye!")
-                return False  # Signal to quit without showing menu
+                maybe_graceful_exit()
+                return False
             else:
                 warn("Please enter 1, 2, or q.")
 

tabpfn_client/tests/quick_test.py

@@ -13,8 +13,8 @@
 from sklearn.model_selection import train_test_split
 
 # from tabpfn_client import UserDataClient
-from tabpfn_client.constants import ModelVersion
 from tabpfn_client.estimator import TabPFNClassifier, TabPFNRegressor
+from tabpfn_client.constants import ModelVersion
 
 logging.basicConfig(level=logging.INFO)
 

tabpfn_client/tests/quick_test_v1.py

@@ -1,6 +1,5 @@
 import os
 import json
-import pandas as pd
 import requests
 from typing import Optional
 
@@ -9,20 +8,22 @@
 TARGET_NAME = "target"
 
 
-def call_fit(train_path: str, target_name: str = TARGET_NAME, api_key: Optional[str] = None) -> str:
+def call_fit(
+    train_path: str, target_name: str = TARGET_NAME, api_key: Optional[str] = None
+) -> str:
     """
     Call the /v1/fit endpoint to train a model.
-    
+
     Args:
         train_path: Path to the training CSV file
         target_name: Name of the target column (default: "target")
         api_key: API key for authentication (if None, reads from PRIORLABS_API_KEY env var)
-    
+
     Returns:
         model_id: The model ID returned from the fit endpoint
     """
     headers = {"Authorization": f"Bearer {api_key}"}
-    
+
     payload = {
         "task": "classification",
         "schema": {
@@ -52,14 +53,14 @@ def call_fit(train_path: str, target_name: str = TARGET_NAME, api_key: Optional[
 def call_predict(test_path: str, model_id: str, api_key: Optional[str] = None) -> None:
     """
     Call the /v1/predict endpoint to get predictions.
-    
+
     Args:
         test_path: Path to the test CSV file
         model_id: The model ID from the fit call
         api_key: API key for authentication (if None, reads from PRIORLABS_API_KEY env var)
     """
     headers = {"Authorization": f"Bearer {api_key}"}
-    
+
     payload = {
         "task": "classification",
         "model_id": model_id,
@@ -99,7 +100,7 @@ def main() -> None:
     # Get the API key
     api_key = os.getenv("PRIORLABS_API_KEY")
 
-    #Test /v1/fit
+    # Test /v1/fit
     print("--- Testing /v1/fit ---")
     model_id = call_fit(train_path, api_key=api_key)
 

tabpfn_client/ui.py

@@ -8,7 +8,6 @@
 from contextlib import contextmanager
 from typing import Iterator, List
 
-from rich import traceback
 from rich.console import Console
 from rich.logging import RichHandler
 from rich.panel import Panel
@@ -33,9 +32,6 @@ def _collect_suppressed_modules() -> List[object]:
     return suppressed
 
 
-traceback.install(show_locals=False, suppress=_collect_suppressed_modules())
-
-
 def _should_use_color() -> bool:
     """Determine whether color output should be used."""