From 290382758bcc7642cf53c87b4151ea332d003a48 Mon Sep 17 00:00:00 2001
From: Dan Lipsitt <578773+DanLipsitt@users.noreply.github.com>
Date: Fri, 25 Jun 2021 18:56:57 -0700
Subject: [PATCH] Fix script to use pip's python version

Problem
-------
On systems where `python` points to `python2` but `pip` points to
`pip3`, `bin/flameprof` breaks because it runs python2 but the
flameprof library is installed into python2's `site-packages`.

Fix
---
Instead of using `scripts` in `setup.py`, use the `console_scripts`
entry
point (https://python-packaging.readthedocs.io/en/latest/command-line-scripts.html). This
will automatically create and install a `flameprof` script that works
with python 2 or 3 (whichever pip is working with).

In order to facilitate this, I created a `main` function in `flameprof.py`.
---
 bin/flameprof | 2 --
 flameprof.py  | 9 +++++----
 setup.py      | 4 +++-
 3 files changed, 8 insertions(+), 7 deletions(-)
 delete mode 100755 bin/flameprof

diff --git a/bin/flameprof b/bin/flameprof
deleted file mode 100755
index 7e83116..0000000
--- a/bin/flameprof
+++ /dev/null
@@ -1,2 +0,0 @@
-#!/bin/sh
-exec python -m flameprof "$@"
diff --git a/flameprof.py b/flameprof.py
index ba9c3a8..db87c9c 100644
--- a/flameprof.py
+++ b/flameprof.py
@@ -371,10 +371,9 @@ def get_arg_parser():
     return parser
 
 
-if __name__ == '__main__':
+def main():
     parser = get_arg_parser()
     args, rest = parser.parse_known_args()
-
     if args.run or args.run_module:
         if args.run:
             code = compile(open(args.stats, mode='rb').read(), '__main__', 'exec', dont_inherit=True)
@@ -402,13 +401,15 @@ def get_arg_parser():
         s.create_stats()
     else:
         s = pstats.Stats(args.stats)
-
     if args.out and args.pstat:
         filename = os.path.splitext(args.out)[0] + '.pstat'
         s.dump_stats(filename)
-
     render(s.stats, get_out(args.out), args.format, args.threshold / 100,
            args.width, args.row_height, args.font_size, args.log_mult)
+
+
+if __name__ == '__main__':
+    main()
 else:
     try:
         import pytest
diff --git a/setup.py b/setup.py
index ced6890..5049be2 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,9 @@
     description='cProfile flamegraph generator',
     long_description=open('README.rst').read(),
     py_modules=['flameprof'],
-    scripts=['bin/flameprof'],
+    entry_points={
+        'console_scripts': ['flameprof=flameprof:main'],
+    },
     include_package_data=True,
     zip_safe=False,
     platforms='any',