add log_time

supercoderhawk · supercoderhawk · commit 92708767bb0b · 2022-05-13T14:29:50.000+08:00
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -2,6 +2,10 @@
 Changelog
 =========
 
+Version 0.1.2
+================
+* add :code:`log_time` method
+
 Version 0.1.1
 ===============
 * fix error in setup.py
diff --git a/pysenal/io/file.py b/pysenal/io/file.py
@@ -54,7 +54,7 @@ def read_lines(filename, encoding=_ENCODING_UTF8, keep_end=False,
 
 
 def read_lines_lazy(filename, encoding=_ENCODING_UTF8, keep_end=False,
-                    strip=False, skip_empty=False, default=None):
+                    strip=False, skip_empty=False, default=None, is_gzip=False):
     """
     use generator to load files, one line every time
     :param filename: source file path
@@ -64,11 +64,16 @@ def read_lines_lazy(filename, encoding=_ENCODING_UTF8, keep_end=False,
     :param skip_empty: whether skip empty line, when strip is False, judge after strip
     :param default: returned value when filename is not existed.
                     If it's None, exception will be raised as usual.
+    :param is_gzip: whether the file is in gzip format
     :return: lines in file one by one
     """
     if not os.path.exists(filename) and default is not None:
         return default
-    file = open(filename, encoding=encoding)
+    if not is_gzip:
+        file = open(filename, encoding=encoding)
+    else:
+        file = gzip.open(filename, 'rt', encoding=encoding)
+        
     for line in file:
         if not keep_end:
             line = line.rstrip(_LINE_BREAKS)
diff --git a/pysenal/utils/logger.py b/pysenal/utils/logger.py
@@ -1,5 +1,6 @@
 # -*- coding: UTF-8 -*-
 import sys
+import time
 import logging
 
 
@@ -13,6 +14,63 @@ def get_logger(name, handler=logging.StreamHandler(sys.stderr), level=logging.DE
     """
     logger = logging.getLogger(name)
     handler.setFormatter(logging.Formatter('[%(asctime)s] [{}] %(message)s'.format(name)))
+    handler.setLevel(level)
     logger.addHandler(handler)
     logger.setLevel(level)
     return logger
+
+
+def log_time(logger=None, log_result=False, log_in_msg=False, request_stage='during'):
+    """
+    func used this decorator must satisfy following condition
+    1. first parameter is dict, indicate the query instance with query_id and start_time
+    2. If log_result is True, logger will print str result or key `log` value of dict result.
+        Result in other data type will raise exception.
+    :param logger: python logger object,
+    :param log_result: whether log func return value
+    :param log_in_msg: whether log time in msg
+    :param request_stage: request stage, default is during
+    :return:
+    """
+    if logger is None:
+        logger = get_logger('default logger')
+
+    def log_time_func(f):
+        new_logger = logger
+
+        def wrap(*args, **kwargs):
+            if args and not isinstance(args[0], dict):
+                if hasattr(args[0], 'logger'):
+                    logger = args[0].logger
+                else:
+                    logger = new_logger
+            else:
+                logger = new_logger
+
+            name = f.__name__
+
+            start = time.time()
+            ret = f(*args, **kwargs)
+            step_time = time.time() - start
+            step_time *= 1000
+
+            info = {'_funcName': name,
+                    'request_stage': request_stage,
+                    'step_time': '{:.0f}ms'.format(step_time)}
+
+            if log_in_msg:
+                ret_str_tmpl = '[func_name:{}] [step_time:{}] [request_stage:{}]'
+                ret_str = ret_str_tmpl.format(name, '{:.0f}ms'.format(step_time), request_stage)
+            else:
+                ret_str = ''
+
+            if log_result:
+                ret_str += ' ' + str(ret)
+                
+            ret_str = ret_str.strip()
+            logger.info(ret_str, extra=info)
+            return ret
+
+        return wrap
+
+    return log_time_func
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 from setuptools import setup, find_packages
 
-VERSION = "0.1.1"
+VERSION = "0.1.2"
 
 requirments = []
 
diff --git a/tests/utils/test_logger.py b/tests/utils/test_logger.py
@@ -1 +1,31 @@
-# -*- coding: UTF-8 -*-
+# -*- coding: UTF-8 -*-
+import logging
+
+from pysenal.utils.logger import *
+
+
+def test_get_logger():
+    logger = get_logger('test_logger')
+    assert logger.name == 'test_logger'
+
+
+def test_log_time():
+    class T(object):
+        def __init__(self):
+            self.logger = get_logger('test_logger', level=logging.INFO)
+
+        @log_time(log_in_msg=True)
+        def test(self):
+            a = 1
+
+        @log_time(log_result=True)
+        def test1(self):
+            return 10
+
+    @log_time(log_in_msg=True)
+    def test_time():
+        a = 1
+
+    assert T().test() is None
+    assert T().test1() == 10
+    assert test_time() is None
