From 850e8e9665eb8821e45ab4c65cb2e831b3fe0b53 Mon Sep 17 00:00:00 2001
From: M Bussonnier <bussonniermatthias@gmail.com>
Date: Wed, 30 Oct 2024 13:30:28 +0100
Subject: [PATCH] more types

---
 lib/python/pyflyby/_interactive.py | 19 +++++++++++++------
 lib/python/pyflyby/_py.py          | 26 +++++++++++++-------------
 lib/python/pyflyby/_util.py        |  3 +--
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/lib/python/pyflyby/_interactive.py b/lib/python/pyflyby/_interactive.py
index 9013fcda..b0c450cc 100644
--- a/lib/python/pyflyby/_interactive.py
+++ b/lib/python/pyflyby/_interactive.py
@@ -15,7 +15,7 @@
 import subprocess
 import sys
 
-from typing import List, Any, Dict
+from typing import List, Any, Dict, Union, Literal
 
 
 from   pyflyby._autoimp         import (LoadSymbolError, ScopeStack, auto_eval,
@@ -2269,7 +2269,9 @@ def disable(self):
 
     def _safe_call(self, function, *args, **kwargs):
         on_error = kwargs.pop("on_error", None)
-        raise_on_error = kwargs.pop("raise_on_error", "if_debug")
+        raise_on_error: Union[bool, Literal["if_debug"]] = kwargs.pop(
+            "raise_on_error", "if_debug"
+        )
         if self._errored:
             # If we previously errored, then we should already have
             # unregistered the hook that led to here.  However, in some corner
@@ -2295,7 +2297,7 @@ def _safe_call(self, function, *args, **kwargs):
                     logger.error("Error trying to disable: %s: %s",
                                  type(e2).__name__, e2)
                 # Raise or print traceback in debug mode.
-                if raise_on_error == True:
+                if raise_on_error is True:
                     raise
                 elif raise_on_error == 'if_debug':
                     if logger.debug_enabled:
@@ -2305,7 +2307,7 @@ def _safe_call(self, function, *args, **kwargs):
                             import traceback
                             traceback.print_exc()
                         raise
-                elif raise_on_error == False:
+                elif raise_on_error is False:
                     if logger.debug_enabled:
                         import traceback
                         traceback.print_exc()
@@ -2328,8 +2330,13 @@ def reset_state_new_cell(self):
                          sorted([k for k,v in autoimported.items() if not v]))
         self._autoimported_this_cell = {}
 
-    def auto_import(self, arg, namespaces=None,
-                    raise_on_error='if_debug', on_error=None):
+    def auto_import(
+        self,
+        arg,
+        namespaces=None,
+        raise_on_error: Union[bool, Literal["if_debug"]] = "if_debug",
+        on_error=None,
+    ):
         if namespaces is None:
             namespaces = get_global_namespaces(self._ip)
 
diff --git a/lib/python/pyflyby/_py.py b/lib/python/pyflyby/_py.py
index 3da73932..834be786 100644
--- a/lib/python/pyflyby/_py.py
+++ b/lib/python/pyflyby/_py.py
@@ -286,19 +286,6 @@
 from   pyflyby._util            import cmp
 from   shlex                    import quote as shquote
 
-usage = """
-py --- command-line python multitool with automatic importing
-
-$ py [--file]   filename.py arg1 arg2      Execute file
-$ py [--apply]  function arg1 arg2         Call function
-$ py [--eval]  'function(arg1, arg2)'      Evaluate code
-$ py [--module] modname arg1 arg2          Run a module
-
-$ py --debug    file/code... args...       Debug code
-$ py --debug    PID                        Attach debugger to PID
-
-$ py                                       IPython shell
-""".strip()
 
 # TODO: add --tidy-imports, etc
 
@@ -377,6 +364,19 @@
 from   pyflyby._parse           import PythonBlock
 from   pyflyby._util            import indent, prefixes
 
+usage = """
+py --- command-line python multitool with automatic importing
+
+$ py [--file]   filename.py arg1 arg2      Execute file
+$ py [--apply]  function arg1 arg2         Call function
+$ py [--eval]  'function(arg1, arg2)'      Evaluate code
+$ py [--module] modname arg1 arg2          Run a module
+
+$ py --debug    file/code... args...       Debug code
+$ py --debug    PID                        Attach debugger to PID
+
+$ py                                       IPython shell
+""".strip()
 
 # Default compiler flags (feature flags) used for all user code.  We include
 # "print_function" here, but we also use auto_flags=True, which means
diff --git a/lib/python/pyflyby/_util.py b/lib/python/pyflyby/_util.py
index d271dded..c525246c 100644
--- a/lib/python/pyflyby/_util.py
+++ b/lib/python/pyflyby/_util.py
@@ -4,7 +4,7 @@
 
 
 
-from   contextlib               import contextmanager
+from   contextlib               import contextmanager, ExitStack
 import inspect
 import os
 import sys
@@ -451,7 +451,6 @@ def cmp(a, b):
 
 
 # Create a context manager with an arbitrary number of contexts.
-from contextlib import ExitStack
 @contextmanager
 def nested(*mgrs):
     with ExitStack() as stack: