Rework run_command so loggers are correctly restored on exit

Pyomo · blnicho · Jan 27, 2021 · Jan 5, 2021 · Jan 18, 2021 · Jan 25, 2021
commit a4e7527caa8a84568295362bd67fb71db7577181
diff --git a/pyomo/scripting/util.py b/pyomo/scripting/util.py
@@ -19,6 +19,8 @@
 import json
 from six import iteritems
 from pyomo.common import pyomo_api
+from pyomo.common.deprecation import deprecated
+from pyomo.common.log import is_debug_set
 from pyomo.common.tempfiles import TempfileManager
 
 from pyutilib.misc import import_file, setup_redirect, reset_redirect
@@ -132,7 +134,7 @@ def pyomo_excepthook(etype,value,tb):
 
         tb_list = traceback.extract_tb(tb,None)
         i = 0
-        if not logger.isEnabledFor(logging.DEBUG) and filter_excepthook:
+        if not is_debug_set(logger) and filter_excepthook:
             while i < len(tb_list):
                 if data.options.model.filename in tb_list[i][0]:
                     break
@@ -414,7 +416,7 @@ def create_model(data):
         total_time = time.time() - modify_start_time
         print("      %6.2f seconds required for problem transformations" % total_time)
 
-    if logger.isEnabledFor(logging.DEBUG):
+    if is_debug_set(logger):
         print("MODEL INSTANCE")
         instance.pprint()
         print("")
@@ -805,51 +807,91 @@ def finalize(data, model=None, instance=None, results=None):
             print('\n# Leaving Interpreter, back to Pyomo\n')
 
 
+@deprecated("configure_loggers is deprecated. The Pyomo command uses the "
+            "PyomoCommandLogContext to update the logger configuration",
+            version='TBD')
 def configure_loggers(options=None, shutdown=False):
+    context = PyomoCommandLogContext(options)
     if shutdown:
-        options = Options()
-        options.runtime = Options()
-        options.runtime.logging = 'quiet'
-        if configure_loggers.fileLogger is not None:
-            logging.getLogger('pyomo').handlers = []
-            logging.getLogger('pyutilib').handlers = []
-            configure_loggers.fileLogger.close()
-            configure_loggers.fileLogger = None
+        # historically, configure_loggers(shutdown=True) forced 'quiet'
+        context.options.runtime.logging = 'quiet'
+        context.fileLogger = configure_loggers.fileLogger
+        context.__exit__()
+    else:
+        context.__enter__()
+        configure_loggers.fileLogger = context.fileLogger
+
+configure_loggers.fileLogger = None
+
+
+class PyomoCommandLogContext(object):
+    """Context manager to setup/restore logging for the Pyomo command"""
+
+    def __init__(self, options):
+        if options is None:
+            options = Options()
+        if options.runtime is None:
+            options.runtime = Options()
+        self.options = options
+        self.fileLogger = None
+        self.original = None
+
+    def __enter__(self):
+        _pyomo = logging.getLogger('pyomo')
+        _pyutilib = logging.getLogger('pyutilib')
+        self.original = ( _pyomo.level, _pyomo.handlers,
+                          _pyutilib.level, _pyutilib.handlers )
+
+        #
+        # Configure the logger
+        #
+        if self.options.runtime.logging == 'quiet':
+            _pyomo.setLevel(logging.ERROR)
+        elif self.options.runtime.logging == 'warning':
+            _pyomo.setLevel(logging.WARNING)
+        elif self.options.runtime.logging == 'info':
+            _pyomo.setLevel(logging.INFO)
+            _pyutilib.setLevel(logging.INFO)
+        elif self.options.runtime.logging == 'verbose':
+            print("VERBOSE!")
+            _pyomo.setLevel(logging.DEBUG)
+            _pyutilib.setLevel(logging.DEBUG)
+        elif self.options.runtime.logging == 'debug':
+            print("DEBUG!")
+            _pyomo.setLevel(logging.DEBUG)
+            _pyutilib.setLevel(logging.DEBUG)
+        elif _pyomo.getEffectiveLevel() == logging.NOTSET:
+            _pyomo.setLevel(logging.WARNING)
+
+        if self.options.runtime.logfile:
+            _logfile = self.options.runtime.logfile
+            self.fileLogger = logging.FileHandler(_logfile, 'w')
+            _pyomo.handlers = []
+            _pyutilib.handlers = []
+            _pyomo.addHandler(self.fileLogger)
+            _pyutilib.addHandler(self.fileLogger)
             # TBD: This seems dangerous in Windows, as the process will
-            # have multiple open file handles pointint to the same file.
-            reset_redirect()
+            # have multiple open file handles pointing to the same file.
+            setup_redirect(_logfile)
 
-    #
-    # Configure the logger
-    #
-    if options.runtime is None:
-        options.runtime = Options()
-    if options.runtime.logging == 'quiet':
-        logging.getLogger('pyomo').setLevel(logging.ERROR)
-    elif options.runtime.logging == 'warning':
-        logging.getLogger('pyomo').setLevel(logging.WARNING)
-    elif options.runtime.logging == 'info':
-        logging.getLogger('pyomo').setLevel(logging.INFO)
-        logging.getLogger('pyutilib').setLevel(logging.INFO)
-    elif options.runtime.logging == 'verbose':
-        logging.getLogger('pyomo').setLevel(logging.DEBUG)
-        logging.getLogger('pyutilib').setLevel(logging.DEBUG)
-    elif options.runtime.logging == 'debug':
-        logging.getLogger('pyomo').setLevel(logging.DEBUG)
-        logging.getLogger('pyutilib').setLevel(logging.DEBUG)
-
-    if options.runtime.logfile:
-        configure_loggers.fileLogger \
-            = logging.FileHandler(options.runtime.logfile, 'w')
-        logging.getLogger('pyomo').handlers = []
-        logging.getLogger('pyutilib').handlers = []
-        logging.getLogger('pyomo').addHandler(configure_loggers.fileLogger)
-        logging.getLogger('pyutilib').addHandler(configure_loggers.fileLogger)
-        # TBD: This seems dangerous in Windows, as the process will
-        # have multiple open file handles pointint to the same file.
-        setup_redirect(options.runtime.logfile)
+        return self
+
+    def __exit__(self, et, ev, tb):
+        _pyomo = logging.getLogger('pyomo')
+        _pyomo.setLevel(self.original[0])
+        _pyomo.handlers = self.original[1]
+        _pyutilib = logging.getLogger('pyutilib')
+        _pyutilib.setLevel(self.original[2])
+        _pyutilib.handlers = self.original[3]
+
+        _pyomo.setLevel(logging.ERROR)
+
+        if self.fileLogger is not None:
+            self.fileLogger.close()
+            # TBD: This seems dangerous in Windows, as the process will
+            # have multiple open file handles pointing to the same file.
+            reset_redirect()
 
-configure_loggers.fileLogger = None
 
 @pyomo_api(namespace='pyomo.script')
 def run_command(command=None, parser=None, args=None, name='unknown', data=None, options=None):
@@ -883,15 +925,14 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
     # Parse command-line options
     #
     #
-    retval = None
-    errorcode = 0
     if options is None:
         try:
             if type(args) is argparse.Namespace:
                 _options = args
             else:
                 _options = parser.parse_args(args=args)
-            # Replace the parser options object with a pyutilib.misc.Options object
+            # Replace the parser options object with a
+            # pyutilib.misc.Options object
             options = Options()
             for key in dir(_options):
                 if key[0] != '_':
@@ -901,15 +942,29 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
         except SystemExit:
             # the parser throws a system exit if "-h" is specified - catch
             # it to exit gracefully.
-            return Container(retval=retval, errorcode=errorcode)
+            return Container(retval=None, errorcode=0)
     #
     # Configure loggers
     #
-    configure_loggers(options=options)
+    TempfileManager.push()
+    try:
+        with PyomoCommandLogContext(options):
+            retval, errorcode = _run_command_impl(
+                command, parser, args, name, data, options)
+    finally:
+        if options.runtime.disable_gc:
+            gc.enable()
+        TempfileManager.pop(remove=not options.runtime.keep_files)
+
+    return Container(retval=retval, errorcode=errorcode)
+
+
+def _run_command_impl(command, parser, args, name, data, options):
     #
     # Call the main Pyomo runner with profiling
     #
-    TempfileManager.push()
+    retval = None
+    errorcode = 0
     pcount = options.runtime.profile_count
     if pcount > 0:
         # Defer import of profiling packages until we know that they
@@ -921,13 +976,13 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
                 import profile
             import pstats
         except ImportError:
-            configure_loggers(shutdown=True)
             raise ValueError(
                 "Cannot use the 'profile' option: the Python "
                 "'profile' or 'pstats' package cannot be imported!")
         tfile = TempfileManager.create_tempfile(suffix=".profile")
         tmp = profile.runctx(
-          command.__name__ + '(options=options,parser=parser)', command.__globals__, locals(), tfile
+            command.__name__ + '(options=options,parser=parser)',
+            command.__globals__, locals(), tfile
         )
         p = pstats.Stats(tfile).strip_dirs()
         p.sort_stats('time', 'cumulative')
@@ -947,7 +1002,6 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
         #
         # Call the main Pyomo runner without profiling
         #
-        TempfileManager.push()
         try:
             retval = command(options=options, parser=parser)
         except SystemExit:
@@ -957,7 +1011,6 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
             # exit.  Otherwise, print an "Exiting..." message.
             #
             if __debug__ and (options.runtime.logging == 'debug' or options.runtime.catch_errors):
-                configure_loggers(shutdown=True)
                 sys.exit(0)
             print('Exiting %s: %s' % (name, str(err)))
             errorcode = err.code
@@ -968,8 +1021,6 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
             # pass the exception up the chain (to pyomo_excepthook)
             #
             if __debug__ and (options.runtime.logging == 'debug' or options.runtime.catch_errors):
-                configure_loggers(shutdown=True)
-                TempfileManager.pop(remove=not options.runtime.keep_files)
                 raise
 
             if not options.model is None and not options.model.save_file is None:
@@ -997,12 +1048,8 @@ def run_command(command=None, parser=None, args=None, name='unknown', data=None,
             logger.error(msg+errStr)
             errorcode = 1
 
-    configure_loggers(shutdown=True)
+    return retval, errorcode
 
-    if options.runtime.disable_gc:
-        gc.enable()
-    TempfileManager.pop(remove=not options.runtime.keep_files)
-    return Container(retval=retval, errorcode=errorcode)
 
 def cleanup():
     for key in modelapi: