Fix line_number and arg_name order and occurrence number (#68)

src/objprint/objprint.py

@@ -70,21 +70,37 @@ def __call__(self, *objs, file=None, format="string", **kwargs):
         cfg = self._configs.overwrite(**kwargs)
         if cfg.enable:
             call_frame = inspect.currentframe().f_back
+
+            # Strip the kwargs that only works in op() so it won't break
+            # json.dumps()
+            kwargs.pop("arg_name", None)
+
+            if cfg.line_number:
+                self._sys_print(self._get_line_number_str(call_frame, cfg=cfg))
+
+            if cfg.arg_name:
+                args = self.frame_analyzer.get_args(call_frame)
+                if args is None:
+                    args = ["Unknown Arg" for _ in range(len(objs))]
+                if cfg.color:
+                    args = [set_color(f"{arg}:", COLOR.RED) for arg in args]
+                else:
+                    args = [f"{arg}:" for arg in args]
+
             if format == "json":
-                for obj in objs:
-                    self._sys_print(json.dumps(self.objjson(obj), **kwargs))
+                if cfg.arg_name:
+                    for arg, obj in zip(args, objs):
+                        self._sys_print(arg)
+                        self._sys_print(json.dumps(self.objjson(obj), **kwargs))
+                else:
+                    for obj in objs:
+                        self._sys_print(json.dumps(self.objjson(obj), **kwargs))
             else:
                 # Force color with cfg as if color is not in cfg, objstr will default to False
                 kwargs["color"] = cfg.color
                 if cfg.arg_name:
-                    args = self.frame_analyzer.get_args(call_frame)
-                    if args is None:
-                        args = ["Unknown Arg" for _ in range(len(objs))]
                     for arg, obj in zip(args, objs):
-                        if cfg.color:
-                            self._sys_print(set_color(f"{arg}:", COLOR.RED))
-                        else:
-                            self._sys_print(f"{arg}:")
+                        self._sys_print(arg)
                         self._sys_print(self.objstr(obj, call_frame=call_frame, **kwargs), file=file)
                 else:
                     for obj in objs:
@@ -102,10 +118,7 @@ def objstr(self, obj, call_frame=None, **kwargs):
             kwargs["color"] = False
         cfg = self._configs.overwrite(**kwargs)
         memo = set() if cfg.skip_recursion else None
-        if cfg.line_number:
-            return self._get_line_number_str(call_frame, cfg=cfg) + self._objstr(obj, memo, indent_level=0, cfg=cfg)
-        else:
-            return self._objstr(obj, memo, indent_level=0, cfg=cfg)
+        return self._objstr(obj, memo, indent_level=0, cfg=cfg)
 
     def _objstr(self, obj, memo, indent_level, cfg):
         # If it's builtin type, return it directly
@@ -229,9 +242,9 @@ def _get_line(key):
     def _get_line_number_str(self, curr_frame, cfg):
         curr_code = curr_frame.f_code
         if cfg.color:
-            return f"{set_color(curr_code.co_name, COLOR.GREEN)} ({curr_code.co_filename}:{curr_frame.f_lineno})\n"
+            return f"{set_color(curr_code.co_name, COLOR.GREEN)} ({curr_code.co_filename}:{curr_frame.f_lineno})"
         else:
-            return f"{curr_code.co_name} ({curr_code.co_filename}:{curr_frame.f_lineno})\n"
+            return f"{curr_code.co_name} ({curr_code.co_filename}:{curr_frame.f_lineno})"
 
     def enable(self):
         self.config(enable=True)

tests/test_objprint.py

@@ -59,6 +59,20 @@ def test_json(self):
             output = buf.getvalue()
         self.assertEqual(output, json.dumps({".type": "ObjTest", "name": "Lisa", "age": 19}, indent=2) + "\n")
 
+    def test_json_with_arg_name(self):
+        with io.StringIO() as buf, redirect_stdout(buf):
+            unique_name = ObjTest({"name": "Lisa", "age": 19})
+            op(unique_name, format="json", arg_name=True)
+            output = buf.getvalue()
+        self.assertIn(json.dumps({".type": "ObjTest", "name": "Lisa", "age": 19}), output)
+        self.assertIn("b", output)
+
+        with io.StringIO() as buf, redirect_stdout(buf):
+            b = ObjTest({"name": "Lisa", "age": 19})
+            op(b, format="json", indent=2)
+            output = buf.getvalue()
+        self.assertEqual(output, json.dumps({".type": "ObjTest", "name": "Lisa", "age": 19}, indent=2) + "\n")
+
     def test_config_indent(self):
         with io.StringIO() as buf, redirect_stdout(buf):
             obj = ObjTest({"name": "Alpha"})

tests/test_objstr.py

@@ -171,11 +171,3 @@ def test_recursion(self):
         self.assertEqual(s.count("t2"), 1)
         s = objstr(t2, skip_recursion=False, depth=6)
         self.assertEqual(s.count("t2"), 3)
-
-    def test_line_number(self):
-        obj = ObjTest({})
-        s = objstr(obj, line_number=True)
-
-        first_line = s.split("\n")[0]
-        self.assertIn("test_line_number", first_line)
-        self.assertIn("test_objstr", first_line)