From 8ca5a871ff3435a5c4dd813d81c2deb12a32225b Mon Sep 17 00:00:00 2001
From: Michal Sojka <michal.sojka@cvut.cz>
Date: Sun, 22 Sep 2024 10:58:27 +0200
Subject: [PATCH] Allow overriding nix-ros-overlay version

---
 README.md          |  9 +++++++--
 ros2nix/ros2nix.py | 40 +++++++++++++++++++++++++++++++++-------
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 2438b2c..2817a42 100644
--- a/README.md
+++ b/README.md
@@ -104,8 +104,8 @@ usage: ros2nix [-h]
                [--extra-check-inputs DEP1,DEP2,...]
                [--extra-native-build-inputs DEP1,DEP2,...] [--flake]
                [--default | --no-default] [--overlay | --no-overlay]
-               [--nixfmt] [--compare] [--copyright-holder COPYRIGHT_HOLDER]
-               [--license LICENSE]
+               [--nix-ros-overlay FLAKEREF] [--nixfmt] [--compare]
+               [--copyright-holder COPYRIGHT_HOLDER] [--license LICENSE]
                package.xml [package.xml ...]
 
 positional arguments:
@@ -165,6 +165,11 @@ options:
                         None)
   --overlay, --no-overlay
                         Generate overlay.nix (default: True)
+  --nix-ros-overlay FLAKEREF
+                        Flake reference of nix-ros-overlay. You may want to
+                        change the branch from master to develop or use your
+                        own fork. (default: github:lopsided98/nix-ros-
+                        overlay/master)
   --nixfmt              Format the resulting expressions with nixfmt (default:
                         False)
   --compare             Don't write any file, but check whether writing the
diff --git a/ros2nix/ros2nix.py b/ros2nix/ros2nix.py
index 4e3a4ec..05adfc3 100755
--- a/ros2nix/ros2nix.py
+++ b/ros2nix/ros2nix.py
@@ -111,17 +111,37 @@ ros_distro_overlays_def = dedent(
 ).strip()
 
 
+def flakeref_to_expr(flakeref) -> str:
+    match flakeref[0]:
+        case '.' | '/':
+            expr = flakeref
+        case _:
+            match re.match("(?P<type>.*?):(?P<owner>.*?)/(?P<repo>.*?)(?:/(?P<ref>.*))?$", flakeref):
+                case None:
+                    raise Exception(f'Unsupported flakeref: "{flakeref}"')
+                case parts:
+                    match parts.groups():
+                        case ('github', owner, repo, None):
+                            expr = f'builtins.fetchTarball "https://github.com/{owner}/{repo}/archive/HEAD.tar.gz"'
+                        case ('github', owner, repo, branch):
+                            expr = f'builtins.fetchTarball "https://github.com/{owner}/{repo}/archive/{branch}.tar.gz"'
+                        case _:
+                            raise Exception(f'Unsupported flakeref: "{flakeref}"')
+    return expr
+
+
 def generate_default(args):
+    nix_ros_overlay = flakeref_to_expr(args.nix_ros_overlay)
     with file_writer(f'{args.output_dir or "."}/default.nix', args.compare) as f:
-        f.write('''{
-  nix-ros-overlay ? builtins.fetchTarball "https://github.com/lopsided98/nix-ros-overlay/archive/master.tar.gz",
-}:
+        f.write(f'''{{
+  nix-ros-overlay ? {nix_ros_overlay},
+}}:
 let
-''' + indent(ros_distro_overlays_def, "  ") + '''
+{indent(ros_distro_overlays_def, "  ")}
 in
-import nix-ros-overlay {
+import nix-ros-overlay {{
   overlays = [ rosDistroOverlays ];
-}
+}}
 ''')
 
 
@@ -130,7 +150,7 @@ def generate_flake(args):
         f.write('''
 {
   inputs = {
-    nix-ros-overlay.url = "github:lopsided98/nix-ros-overlay/master";
+    nix-ros-overlay.url = "''' + args.nix_ros_overlay + '''";
     nixpkgs.follows = "nix-ros-overlay/nixpkgs";  # IMPORTANT!!!
   };
   outputs = { self, nix-ros-overlay, nixpkgs }:
@@ -266,6 +286,12 @@ def ros2nix(args):
         default=True,
         help="Generate overlay.nix",
     )
+    parser.add_argument(
+        "--nix-ros-overlay",
+        metavar="FLAKEREF",
+        default="github:lopsided98/nix-ros-overlay/master",
+        help="Flake reference of nix-ros-overlay. You may want to change the branch from master to develop or use your own fork.",
+    )
     parser.add_argument(
         "--nixfmt",
         action="store_true",