Add profile argument to update and claim commands

Author: Douglas Lassance

gitalong/batch.py

@@ -413,7 +413,6 @@ async def update_tracked_commits(
     )
     absolute_filenames = []
     if repository.config.get("modify_permissions"):
-        print("Updating permissions")
         for filename in repository.files:
             absolute_filenames.append(repository.get_absolute_path(filename))
         await repository.batch.update_files_permissions(absolute_filenames)

gitalong/cli.py

@@ -7,6 +7,7 @@
 from .__info__ import __version__
 from .enums import CommitSpread
 from .repository import Repository
+from .contexts import ProfileContext
 from .batch import get_files_last_commits, claim_files, update_tracked_commits
 
 
@@ -66,9 +67,25 @@ def config(ctx, prop):  # pylint: disable=missing-function-docstring
         "files their permissions changed."
     )
 )
+@click.option(
+    "-p",
+    "--profile",
+    is_flag=True,
+    help=(
+        "Will generate a profile file in the current working directory."
+        "The file can be opened in a profiler like snakeviz."
+    ),
+)
 @click.pass_context
-def update(ctx):
-    """Update tracked commits with local changes."""
+def update(ctx, profile=False):  # pylint: disable=missing-function-docstring
+    if profile:
+        with ProfileContext():
+            run_update(ctx)
+        return
+    run_update(ctx)
+
+
+def run_update(ctx):  # pylint: disable=missing-function-docstring
     repository = Repository.from_filename(ctx.obj.get("REPOSITORY", ""))
     if repository:
         asyncio.run(update_tracked_commits(repository))
@@ -94,13 +111,8 @@ def update(ctx):
 @click.pass_context
 def status(ctx, filename, profile=False):  # pylint: disable=missing-function-docstring
     if profile:
-        import cProfile  # pylint: disable=import-outside-toplevel
-        import pstats  # pylint: disable=import-outside-toplevel
-
-        with cProfile.Profile() as pr:
+        with ProfileContext():
             run_status(ctx, filename)
-        results = pstats.Stats(pr)
-        results.dump_stats("gitalong.prof")
         return
     run_status(ctx, filename)
 
@@ -132,8 +144,25 @@ def run_status(ctx, filename):  # pylint: disable=missing-function-docstring
     nargs=-1,
     # help="The path to the file that should be made writable."
 )
+@click.option(
+    "-p",
+    "--profile",
+    is_flag=True,
+    help=(
+        "Will generate a profile file in the current working directory."
+        "The file can be opened in a profiler like snakeviz."
+    ),
+)
 @click.pass_context
-def claim(ctx, filename):  # pylint: disable=missing-function-docstring
+def claim(ctx, filename, profile=False):  # pylint: disable=missing-function-docstring
+    if profile:
+        with ProfileContext():
+            run_claim(ctx, filename)
+        return
+    run_claim(ctx, filename)
+
+
+def run_claim(ctx, filename):  # pylint: disable=missing-function-docstring
     absolute_filenames = []
     for filename_ in filename:
         repository = Repository.from_filename(ctx.obj.get("REPOSITORY", filename_))

gitalong/commit.py

@@ -148,7 +148,7 @@ def commit_spread(self) -> int:
                 information about where this commit lives across branches and clones.
         """
         commit_spread = 0
-        active_branch = self._repository.active_branch_name
+        active_branch = self._repository.active_branch_name if self._repository else ""
         if self.get("user", ""):
             is_issued = self.is_issued_commit()
             if "sha" in self:

gitalong/contexts.py

@@ -0,0 +1,19 @@
+import cProfile
+import pstats
+
+
+class ProfileContext:
+    """Profile context manager."""
+
+    def __init__(self):
+        self._profile = None
+
+    def __enter__(self):
+        self._profile = cProfile.Profile()
+        self._profile.enable()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._profile:
+            self._profile.disable()
+            results = pstats.Stats(self._profile)
+            results.dump_stats("gitalong.prof")