unify logging level setup

nbcsm · nbcsm · commit 8b0d5cd58d3e · 2019-04-16T11:32:54.000+08:00
diff --git a/tests/common.py b/tests/common.py
@@ -4,6 +4,7 @@
 """ test common utilities."""
 
 import argparse
+import logging
 import os
 import sys
 import unittest
@@ -114,7 +115,10 @@ def get_test_config():
 
 
 def unittest_main():
-    print(get_test_config())
+    config = get_test_config()
+    logging.basicConfig(level=logging.WARNING)
+    with utils.set_log_level(logging.getLogger(), logging.INFO) as logger:
+        logger.info(config)
     unittest.main()
 
 
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -3,8 +3,13 @@
 
 """ print pytest config."""
 
+import logging
 from common import get_test_config
+from tf2onnx import utils
 
 
 def pytest_configure():
-    print(get_test_config())
+    config = get_test_config()
+    logging.basicConfig(level=logging.WARNING)
+    with utils.set_log_level(logging.getLogger(), logging.INFO) as logger:
+        logger.info(config)
diff --git a/tf2onnx/convert.py b/tf2onnx/convert.py
@@ -10,6 +10,7 @@
 from __future__ import unicode_literals
 
 import argparse
+import logging
 import tensorflow as tf
 
 from tf2onnx import constants, loader, utils
@@ -79,6 +80,7 @@ def default_custom_op_handler(ctx, node, name, args):
 
 def main():
     args = get_args()
+    logging.basicConfig(level=logging.INFO)
 
     # override unknown dimensions from -1 to 1 (aka batchsize 1) since not every runtime does
     # support unknown dimensions.
diff --git a/tf2onnx/graph.py b/tf2onnx/graph.py
@@ -23,7 +23,7 @@
 from tf2onnx import optimizer
 from tf2onnx.schemas import get_schema
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("graph")
 
 
diff --git a/tf2onnx/loader.py b/tf2onnx/loader.py
@@ -15,7 +15,7 @@
 
 from tf2onnx import utils
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("loader")
 
 # pylint: disable=unused-argument
diff --git a/tf2onnx/onnx_opset/common.py b/tf2onnx/onnx_opset/common.py
@@ -13,7 +13,7 @@
 
 from tf2onnx import constants
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.common")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/controlflow.py b/tf2onnx/onnx_opset/controlflow.py
@@ -20,7 +20,7 @@
 from tf2onnx.handler import tf_op
 from tf2onnx.utils import make_sure
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.controlflow")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/generator.py b/tf2onnx/onnx_opset/generator.py
@@ -16,7 +16,7 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.generator")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/logical.py b/tf2onnx/onnx_opset/logical.py
@@ -16,7 +16,7 @@
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.logical")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/math.py b/tf2onnx/onnx_opset/math.py
@@ -17,7 +17,7 @@
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.math")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/misc.py b/tf2onnx/onnx_opset/misc.py
@@ -13,7 +13,7 @@
 
 from tf2onnx.handler import tf_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.misc")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/nn.py b/tf2onnx/onnx_opset/nn.py
@@ -18,7 +18,7 @@
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common, controlflow, tensor
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.nn")
 
 # pylint: disable=unused-argument,missing-docstring,unused-variable
diff --git a/tf2onnx/onnx_opset/reduction.py b/tf2onnx/onnx_opset/reduction.py
@@ -17,7 +17,7 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.reduction")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/onnx_opset/rnn.py b/tf2onnx/onnx_opset/rnn.py
@@ -15,7 +15,7 @@
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.rnn")
 
 
diff --git a/tf2onnx/onnx_opset/tensor.py b/tf2onnx/onnx_opset/tensor.py
@@ -19,7 +19,7 @@
 from tf2onnx import constants, utils
 from tf2onnx.handler import tf_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.tensor")
 
 # pylint: disable=unused-argument,missing-docstring,unused-variable
diff --git a/tf2onnx/onnx_opset/traditionalml.py b/tf2onnx/onnx_opset/traditionalml.py
@@ -11,7 +11,7 @@
 
 import logging
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("onnx_opset.traditional")
 
 # pylint: disable=unused-argument,missing-docstring
diff --git a/tf2onnx/rewriter/bigru_rewriter.py b/tf2onnx/rewriter/bigru_rewriter.py
@@ -18,7 +18,7 @@
      get_reverse_nodes_after_y_output, get_np_val_for_const, _process_single_init_node
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.bigru_rewriter")
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
diff --git a/tf2onnx/rewriter/bilstm_rewriter.py b/tf2onnx/rewriter/bilstm_rewriter.py
@@ -15,7 +15,7 @@
 from tf2onnx import utils
 from tf2onnx.rewriter.rnn_utils import is_reverse_op
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.bilstm_rewriter")
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
diff --git a/tf2onnx/rewriter/cond_rewriter.py b/tf2onnx/rewriter/cond_rewriter.py
@@ -13,7 +13,7 @@
 from enum import Enum
 from tf2onnx import utils
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.cond_rewriter_base")
 
 
diff --git a/tf2onnx/rewriter/custom_rnn_rewriter.py b/tf2onnx/rewriter/custom_rnn_rewriter.py
@@ -17,7 +17,7 @@
 from tf2onnx.tfonnx import utils
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.custom_rnn_rewriter")
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test,broad-except,protected-access
diff --git a/tf2onnx/rewriter/gru_rewriter.py b/tf2onnx/rewriter/gru_rewriter.py
@@ -18,7 +18,7 @@
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.gru_rewriter")
 
 
diff --git a/tf2onnx/rewriter/loop_rewriter.py b/tf2onnx/rewriter/loop_rewriter.py
@@ -17,7 +17,7 @@
 from tf2onnx.tfonnx import utils
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.loop_rewriter")
 
 
diff --git a/tf2onnx/rewriter/loop_rewriter_base.py b/tf2onnx/rewriter/loop_rewriter_base.py
@@ -17,7 +17,7 @@
 from tf2onnx.rewriter.rnn_utils import REWRITER_RESULT
 from tf2onnx.utils import TensorValueInfo
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.loop_rewriter_base")
 INVALID_INPUT_ID = utils.make_name("invalid_input_id")
 
diff --git a/tf2onnx/rewriter/lstm_rewriter.py b/tf2onnx/rewriter/lstm_rewriter.py
@@ -19,7 +19,7 @@
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.lstm_rewriter")
 
 
diff --git a/tf2onnx/rewriter/rnn.py b/tf2onnx/rewriter/rnn.py
@@ -20,7 +20,7 @@
 
 # pylint: disable=invalid-name,unused-argument,missing-docstring
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.rnn")
 
 
diff --git a/tf2onnx/rewriter/rnn_utils.py b/tf2onnx/rewriter/rnn_utils.py
@@ -16,7 +16,7 @@
 # pylint: disable=invalid-name,unused-argument,missing-docstring
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.rnn_utils")
 
 
diff --git a/tf2onnx/rewriter/unit_rnn_rewriter_base.py b/tf2onnx/rewriter/unit_rnn_rewriter_base.py
@@ -16,7 +16,7 @@
 from tf2onnx.graph_matcher import GraphMatcher
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.rewriter.unit_rnn_rewriter_base")
 
 
diff --git a/tf2onnx/shape_inference.py b/tf2onnx/shape_inference.py
@@ -15,7 +15,7 @@
 # pylint: disable=logging-not-lazy,missing-docstring,consider-swap-variables
 
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx.shape_inference")
 
 direct_ops = [
diff --git a/tf2onnx/tfonnx.py b/tf2onnx/tfonnx.py
@@ -30,7 +30,7 @@
 from tf2onnx.shape_inference import infer_shape_for_graph
 from tf2onnx.utils import port_name
 
-logging.basicConfig(level=logging.INFO)
+
 log = logging.getLogger("tf2onnx")
 
 
diff --git a/tf2onnx/utils.py b/tf2onnx/utils.py
@@ -9,6 +9,7 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+from contextlib import contextmanager
 import os
 import re
 import shutil
@@ -437,3 +438,15 @@ def is_onnx_domain(domain):
     if domain is None or domain == "":
         return True
     return False
+
+
+@contextmanager
+def set_log_level(logger, level):
+    """ Override current logger level within context. """
+    current_level = logger.getEffectiveLevel()
+    logger.setLevel(level)
+
+    try:
+        yield logger
+    finally:
+        logger.setLevel(current_level)
