From 958c077da36c4e5b488122bfd74306254b6fc70d Mon Sep 17 00:00:00 2001
From: Jan Lindemann <jan@janware.com>
Date: Thu, 5 Jun 2025 16:58:26 +0200
Subject: [PATCH] Cmds: Support --write-profile <filename>

Add an option which makes the program write profiling statistics in pstats
format to a specified file.

Visualize e.g. with:

  gprof2dot -f pstats <filename> -o <dotfile>
  dot -Tpng | display <dotfile>

Signed-off-by: Jan Lindemann <jan@janware.com>
---
 tools/python/jwutils/Cmds.py | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/tools/python/jwutils/Cmds.py b/tools/python/jwutils/Cmds.py
index 5385252..6f71705 100644
--- a/tools/python/jwutils/Cmds.py
+++ b/tools/python/jwutils/Cmds.py
@@ -9,6 +9,8 @@ import asyncio
 from argparse import ArgumentParser
 from typing import Optional
 
+import cProfile
+
 import jwutils
 from jwutils.log import *
 
@@ -68,6 +70,7 @@ class Cmds: # export
         self.__parser.add_argument('--log-flags', help='Log flags', default=log_flags)
         self.__parser.add_argument('--log-level', help='Log level', default=log_level)
         self.__parser.add_argument('--backtrace', help='Show exception backtraces', action='store_true', default=False)
+        self.__parser.add_argument('--write-profile', help='Profile code and store output to file', default=None)
         if self.__modules == None:
             self.__modules = [ '__main__' ]
         subcmds = set()
@@ -99,6 +102,7 @@ class Cmds: # export
         set_flags(self.args.log_flags)
         set_level(self.args.log_level)
         self.__back_trace = self.args.backtrace
+        exit_status = 0
 
         # This is the toplevel parser, i.e. no func member has been added to the args via
         #
@@ -112,17 +116,27 @@ class Cmds: # export
             self.__parser.print_help()
             return None
 
-        if self.__back_trace:
+        pr = None if self.args.write_profile is None else cProfile.Profile()
+        if pr is not None:
+            pr.enable()
+
+        try:
             ret = await self.args.func(self.args)
-        else:
-            try:
-                ret = await self.args.func(self.args)
-            except Exception as e:
-                if hasattr(e, 'message'):
-                    slog(ERR, e.message)
-                else:
-                    slog(ERR, f'Exception: {type(e)}: {e}')
-                sys.exit(1)
+        except Exception as e:
+            if hasattr(e, 'message'):
+                slog(ERR, e.message)
+            else:
+                slog(ERR, f'Exception: {type(e)}: {e}')
+            exit_status = 1
+            if self.__back_trace:
+                raise
+        finally:
+            if pr is not None:
+                pr.disable()
+                slog(NOTICE, f'Writing profile statistics to {self.args.write_profile}')
+                pr.dump_stats(self.args.write_profile)
+
+        sys.exit(exit_status)
 
     def __del__(self):
         if self.__own_eloop: