5252import json
5353import os
5454import re
55+ from functools import cache
5556from subprocess import SubprocessError
5657
5758from stack_pr .git import (
@@ -382,12 +383,14 @@ def split_header(s: str) -> List[CommitHeader]:
382383 return [CommitHeader (h ) for h in s .split ("\0 " )[:- 1 ]]
383384
384385
385- def is_valid_ref (ref : str ) -> bool :
386+ def is_valid_ref (ref : str , branch_name_template : str ) -> bool :
386387 ref = ref .strip ("'" )
387- splits = ref .rsplit ("/" , 2 )
388- if len (splits ) < 3 :
388+
389+ branch_name_base = get_branch_name_base (branch_name_template )
390+ splits = ref .rsplit ("/" , 1 )
391+ if len (splits ) < 2 :
389392 return False
390- return splits [- 2 ] == "stack" and splits [- 1 ].isnumeric ()
393+ return splits [- 2 ]. endswith ( branch_name_base ) and splits [- 1 ].isnumeric ()
391394
392395
393396def last (ref : str , sep : str = "/" ) -> str :
@@ -555,43 +558,57 @@ def add_or_update_metadata(
555558 return True
556559
557560
558- def get_available_branch_name (remote : str ) -> str :
561+ @cache
562+ def get_branch_name_base (branch_name_template : str ):
559563 username = get_gh_username ()
564+ branch_name_base = branch_name_template .replace ("$USERNAME" , username )
565+ return branch_name_base
566+
567+
568+ def get_available_branch_name (remote : str , branch_name_template : str ) -> str :
569+ branch_name_base = get_branch_name_base (branch_name_template )
560570
561571 refs = get_command_output (
562572 [
563573 "git" ,
564574 "for-each-ref" ,
565- f"refs/remotes/{ remote } /{ username } /stack " ,
575+ f"refs/remotes/{ remote } /{ branch_name_base } " ,
566576 "--format='%(refname)'" ,
567577 ]
568578 ).split ()
569579
570- refs = list (filter (is_valid_ref , refs ))
580+ def check_ref (ref ):
581+ return is_valid_ref (ref , branch_name_base )
582+
583+ refs = list (filter (check_ref , refs ))
571584 max_ref_num = max (int (last (ref .strip ("'" ))) for ref in refs ) if refs else 0
572585 new_branch_id = max_ref_num + 1
573586
574- return f"{ username } /stack /{ new_branch_id } "
587+ return f"{ branch_name_base } /{ new_branch_id } "
575588
576589
577590def get_next_available_branch_name (name : str ) -> str :
578591 base , id = name .rsplit ("/" , 1 )
579592 return f"{ base } /{ int (id ) + 1 } "
580593
581594
582- def set_head_branches (st : List [StackEntry ], remote : str , verbose : bool ):
595+ def set_head_branches (
596+ st : List [StackEntry ], remote : str , verbose : bool , branch_name_template : str
597+ ):
583598 """Set the head ref for each stack entry if it doesn't already have one."""
584599
585600 run_shell_command (["git" , "fetch" , "--prune" , remote ], quiet = not verbose )
586- available_name = get_available_branch_name (remote )
601+ available_name = get_available_branch_name (remote , branch_name_template )
587602 for e in filter (lambda e : not e .has_head (), st ):
588603 e .head = available_name
589604 available_name = get_next_available_branch_name (available_name )
590605
591606
592- def init_local_branches (st : List [StackEntry ], remote : str , verbose : bool ):
607+ def init_local_branches (
608+ st : List [StackEntry ], remote : str , verbose : bool , branch_name_template : str
609+ ):
593610 log (h ("Initializing local branches" ), level = 1 )
594- set_head_branches (st , remote , verbose )
611+ set_head_branches (st , remote , verbose , branch_name_template )
595612 for e in st :
596613 run_shell_command (
597614 ["git" , "checkout" , e .commit .commit_id (), "-B" , e .head ],
@@ -785,6 +802,7 @@ class CommonArgs(NamedTuple):
785802 target : str
786803 hyperlinks : bool
787804 verbose : bool
805+ branch_name_template : str
788806
789807 @classmethod
790808 def from_args (cls , args : argparse .Namespace ) -> "CommonArgs" :
@@ -795,6 +813,7 @@ def from_args(cls, args: argparse.Namespace) -> "CommonArgs":
795813 args .target ,
796814 args .hyperlinks ,
797815 args .verbose ,
816+ args .branch_name_template ,
798817 )
799818
800819
@@ -822,6 +841,7 @@ def deduce_base(args: CommonArgs) -> CommonArgs:
822841 args .target ,
823842 args .hyperlinks ,
824843 args .verbose ,
844+ args .branch_name_template ,
825845 )
826846
827847
@@ -876,7 +896,9 @@ def command_submit(
876896
877897 # Create local branches and initialize base and head fields in the stack
878898 # elements
879- init_local_branches (st , args .remote , args .verbose )
899+ init_local_branches (
900+ st , args .remote , args .verbose , args .branch_name_template
901+ )
880902 set_base_branches (st , args .target )
881903 print_stack (st , args .hyperlinks )
882904
@@ -1137,7 +1159,9 @@ def command_abandon(args: CommonArgs):
11371159 return
11381160 current_branch = get_current_branch_name ()
11391161
1140- init_local_branches (st , args .remote , args .verbose )
1162+ init_local_branches (
1163+ st , args .remote , args .verbose , args .branch_name_template
1164+ )
11411165 set_base_branches (st , args .target )
11421166 print_stack (st , args .hyperlinks )
11431167
@@ -1219,7 +1243,7 @@ def command_view(args: CommonArgs):
12191243
12201244 st = get_stack (args .base , args .head , args .verbose )
12211245
1222- set_head_branches (st , args .remote , args .verbose )
1246+ set_head_branches (st , args .remote , args .verbose , args . branch_name_template )
12231247 set_base_branches (st , args .target )
12241248 print_stack (st , args .hyperlinks )
12251249 print_tips_after_view (st , args )
@@ -1268,6 +1292,13 @@ def create_argparser(
12681292 default = config .getboolean ("common" , "verbose" , fallback = False ),
12691293 help = "Enable verbose output from Git subcommands." ,
12701294 )
1295+ common_parser .add_argument (
1296+ "--branch-name-template" ,
1297+ default = config .get (
1298+ "repo" , "branch_name_template" , fallback = "$USERNAME/stack"
1299+ ),
1300+ help = "A template for names of the branches stack-pr would use." ,
1301+ )
12711302
12721303 parser_submit = subparsers .add_parser (
12731304 "submit" ,
0 commit comments