feat: Support nixpkgs URIs

bin/main.ml

@@ -7,34 +7,34 @@ let main
         Cli.strategy)
   =
   let ({ entrypoint; attribute; subshell_dir } : Cli.target_info) = target_info in
-  Option.value ~default:(Option.value ~default:(Sys.getcwd ()) entrypoint) subshell_dir
-  |> Unix.cd;
-  let cmd =
-    match entrypoint with
-    | Some entrypoint ->
-      (match Unix.flake_exists_at entrypoint with
-       | true ->
-         Cmd.builder "nix"
-         |>+ [ "develop" ] @ [ Uri.sprintf_uri_attr_opt entrypoint attribute ]
-         |>+ Option.value ~default:[] force_experimental_features
-       | false ->
-         (match Unix.shell_exists_at entrypoint with
-          | true ->
-            Cmd.builder "nix-shell"
-            |>+ Option.value
-                  ~default:[]
-                  (Option.map (fun attr -> [ "--attr"; attr ]) attribute)
-            |>+ [ entrypoint ]
-          | false ->
-            Error.handle_ns_error "no available devshell entrypoint: %s\n%!" entrypoint))
-      |>+ [ "--command"; Unix.shell ]
-    | None ->
-      Cmd.builder "nix"
-      |>+ [ "shell" ]
-      |>+ installables
-      |>+ Option.value ~default:[] force_experimental_features
+  let strategy =
+    let workdir =
+      Option.value
+        ~default:(Option.value ~default:(Sys.getcwd ()) entrypoint)
+        subshell_dir
+    and primary, fallback =
+      match entrypoint with
+      | Some entrypoint ->
+        if Unix.flake_exists_at entrypoint
+        then
+          ( Cmd.nix_develop entrypoint attribute force_experimental_features
+          , if Unix.shell_exists_at entrypoint
+            then Some (Cmd.legacy_nix_shell_from_entrypoint entrypoint attribute)
+            else None )
+        else if Unix.shell_exists_at entrypoint
+        then Cmd.legacy_nix_shell_from_entrypoint entrypoint attribute, None
+        else Error.handle_ns_error "no available devshell entrypoint: %s\n%!" entrypoint
+      | None ->
+        ( Cmd.nix_shell
+            (List.map Uri.uri_to_string installables)
+            force_experimental_features
+        , Option.map
+            (fun installables -> Cmd.legacy_nix_shell_from_installables installables)
+            (Uri.combine_installables_tr installables) )
+    in
+    { workdir; primary; fallback }
   in
-  if printcmd then print_endline (Cmd.to_string cmd) else ignore (Cmd.run cmd)
+  if printcmd then print_strategy strategy else execute_strategy strategy
 ;;
 
 Cli.eval main

lib/cli.ml

@@ -15,7 +15,7 @@ module Cli = struct
   (** Processed args that will be consumed by the main function to determine the
   behavior of the program. *)
   type strategy =
-    { installables : string list
+    { installables : Uri.uri list
     ; target_info : target_info
     ; printcmd : bool
     ; force_experimental_features : string list option
@@ -25,7 +25,12 @@ module Cli = struct
   let uris = Arg.(value & pos_all dirpath [] & info [])
 
   let printcmd =
-    let doc = "Print the command to stdout instead of executing it." in
+    let doc =
+      "Print the command strategy as newline separated values to stdout instead of \
+       executing it. The first command changes to the working directory, the second is \
+       the primary nix command and the last is an optional fallback nix command in the \
+       event that flake/nix-command features fail."
+    in
     Arg.(value & flag & info [ "printcmd" ] ~doc)
   ;;
 
@@ -64,7 +69,9 @@ module Cli = struct
         (* Two args passed *)
         (*******************)
         | [ maybe_subshell_dir; target ] ->
-          (match Uri.parse_target maybe_subshell_dir, Uri.parse_target target with
+          let maybe_subshell_dir = Uri.parse_target maybe_subshell_dir
+          and target = Uri.parse_target target in
+          (match maybe_subshell_dir, target with
            (* If maybe_subshell_dir does not have any attributes we know the user wants to change directories. *)
            | ( LocalResourceMaybeAttr (subshell_dir, None)
              , LocalResourceMaybeAttr (entrypoint, attribute) ) ->
@@ -73,12 +80,10 @@ module Cli = struct
                ; attribute
                ; subshell_dir = Some subshell_dir
                } )
-           | LocalResourceMaybeAttr (subshell_dir, None), LocalResourceMultiAttr _ ->
+           | ( LocalResourceMaybeAttr (subshell_dir, None)
+             , (LocalResourceMultiAttr _ | RemoteResource _) ) ->
              ( [ target ]
              , { entrypoint = None; attribute = None; subshell_dir = Some subshell_dir } )
-           | LocalResourceMaybeAttr (subshell_dir, None), RemoteResource uri ->
-             ( [ uri ]
-             , { entrypoint = None; attribute = None; subshell_dir = Some subshell_dir } )
            | _ -> Uri.parse_targets_tr original_args, default_target)
         (*****************************)
         (* Three or more args passed *)

lib/cmd.ml

@@ -1,28 +1,85 @@
 (** Functions and operators used in running commands *)
 module Cmd = struct
-  open Bos
+  open Util.Util
 
-  (** Construct a command builder
-
-    Example:
-    Cmd.builder "nix" |>+ \[ "flake" "show" "--json" \]
-  *)
-  let builder cmd = Cmd.v cmd
+  type strategy =
+    { workdir : string
+    ; primary : Bos.Cmd.t
+    ; fallback : Bos.Cmd.t option
+    }
 
   (** Pipe-arg operator, used in adding arguments to a command *)
   let ( |>+ ) cmd = function
     | [] -> cmd
-    | [ arg ] -> Cmd.add_arg cmd arg
-    | args -> List.fold_left (fun cmd arg -> Cmd.add_arg cmd arg) cmd args
+    | [ arg ] -> Bos.Cmd.add_arg cmd arg
+    | args -> List.fold_left (fun cmd arg -> Bos.Cmd.add_arg cmd arg) cmd args
+  ;;
+
+  let nix_develop entrypoint attribute force_experimental_features =
+    Bos.Cmd.v "nix"
+    |>+ [ "develop" ] @ [ Uri.sprintf_uri_attr_opt entrypoint attribute ]
+    |>+ Option.value ~default:[] force_experimental_features
+    |>+ [ "--command"; Unix.shell ]
+  ;;
+
+  let legacy_nix_shell_from_entrypoint entrypoint attribute =
+    Bos.Cmd.v "nix-shell"
+    |>+ Option.value ~default:[] (Option.map (fun attr -> [ "--attr"; attr ]) attribute)
+    |>+ [ entrypoint ]
+    |>+ [ "--command"; Unix.shell ]
   ;;
 
-  (** Convert a command into a string *)
-  let to_string cmd = Cmd.to_string cmd
+  let nix_shell installables force_experimental_features =
+    Bos.Cmd.v "nix"
+    |>+ [ "shell" ]
+    |>+ installables
+    |>+ Option.value ~default:[] force_experimental_features
+  ;;
 
-  (** Run a command
+  let legacy_nix_shell_from_installables installables =
+    (* Need to combine/validate that the installables given are all Nixpkgs *)
+    Bos.Cmd.v "nix-shell"
+    |>+ [ "--packages" ] @ installables
+    |>+ [ "--command"; Unix.shell ]
+  ;;
 
-    Example:
-    Cmd.run (Cmd.builder "nix" |>+ \[ "show-config" "--json" \])
-  *)
-  let run cmd = OS.Cmd.run cmd
+  let print_strategy ({ workdir; primary; fallback } : strategy) =
+    print_string
+      (Printf.sprintf
+         "'cd' '%s'\n%s\n%s"
+         workdir
+         (Bos.Cmd.to_string primary)
+         (Option.value ~default:"" (Option.map Bos.Cmd.to_string fallback)))
+  ;;
+
+  let execute_strategy ({ workdir; primary; fallback } : strategy) =
+    Unix.cd workdir;
+    match Bos.OS.Cmd.run_status primary with
+    | Ok (`Exited 0) -> () (* success, nothing else to do *)
+    | Ok (`Exited code | `Signaled code) ->
+      (* primary failed *)
+      (match fallback with
+       | Some fallback ->
+         (* show primary error but don't exit *)
+         Printf.eprintf "primary command failed (exit %d), trying fallback\n%!" code;
+         (match Bos.OS.Cmd.run_status fallback with
+          | Ok (`Exited 0) -> ()
+          | Ok (`Exited code | `Signaled code) -> exit code
+          | Error (`Msg msg) ->
+            prerr_endline msg;
+            exit 1)
+       | None -> exit code)
+    | Error (`Msg msg) ->
+      (* primary couldn't even be spawned *)
+      prerr_endline msg;
+      (match fallback with
+       | Some fallback ->
+         (match Bos.OS.Cmd.run_status fallback with
+          | Ok (`Exited 0) -> ()
+          | Ok (`Exited code | `Signaled code) -> exit code
+          | Error (`Msg msg) ->
+            prerr_endline msg;
+            exit 1)
+       | None -> exit 1)
+  ;;
 end

lib/util.ml

@@ -40,6 +40,12 @@ module Util = struct
   module Uri = struct
     (* *)
 
+    (** A local reference to nixpkgs, eg. {nixpkgs#...} or {pkgs#...} (which resolves to the former) *)
+    type nixpkgs =
+      { uri : string
+      ; installables : string list
+      }
+
     (** URI variants, like {/path/to/flake#drv} or {github:NixOS/nixpkgs#drv} *)
     type uri =
       (* A path with up to one attribute, eg. /path/to/resouce or /path/to/resource#... *)
@@ -48,6 +54,8 @@ module Util = struct
       | LocalResourceMultiAttr of string * string
       (* A remote resource, eg. github:NixOS/nixpkgs *)
       | RemoteResource of string
+      (* A local reference to nixpkgs, eg. nixpkgs#... or pkgs#... (which resolves to the former) *)
+      | Nixpkgs of nixpkgs
 
     (** Format a uri and attribute into a string *)
     let sprintf_uri_attr path attr = Printf.sprintf "%s#%s" path attr
@@ -65,17 +73,38 @@ module Util = struct
       | LocalResourceMaybeAttr (path, attr_opt) -> sprintf_uri_attr_opt path attr_opt
       | LocalResourceMultiAttr (path, attr) -> sprintf_uri_attr path attr
       | RemoteResource uri -> uri
+      | Nixpkgs { uri; _ } -> uri
     ;;
 
     (** Parse target argument for optional attribute selection *)
     let parse_target target =
       let parse_uri uri attr_opt =
+        let is_multiattr maybe_multiattr =
+          String.starts_with ~prefix:"{" maybe_multiattr
+          && String.ends_with ~suffix:"}" maybe_multiattr
+        in
         match String.split_on_char ':' uri with
-        | [ path ] ->
-          let path = Unix.realpath path
-          and is_multiattr maybe_multiattr =
-            String.starts_with ~prefix:"{" maybe_multiattr
+        | [ "nixpkgs" ] | [ "pkgs" ] ->
+          let parse_installables attr_opt =
+            match attr_opt with
+            | Some attr ->
+              if is_multiattr attr
+              then (
+                match String.split_on_char '{' attr with
+                | [ _; rhs ] ->
+                  (match String.split_on_char '}' rhs with
+                   | [ lhs; _ ] -> String.split_on_char ',' lhs
+                   | _ -> [])
+                | _ -> [])
+              else [ attr ]
+            | None -> []
           in
+          Nixpkgs
+            { uri = sprintf_uri_attr_opt "nixpkgs" attr_opt
+            ; installables = parse_installables attr_opt
+            }
+        | [ path ] ->
+          let path = Unix.realpath path in
           (match attr_opt with
            | Some attr ->
              if is_multiattr attr
@@ -91,15 +120,24 @@ module Util = struct
         Error.handle_ns_error "invalid uri or attribute selection syntax: %s\n%!" target
     ;;
 
-    (** Parse target arguments recursively, returning a list of strings.
+    (** Parse target arguments recursively, returning a list of {!type:uri}.
     This function is tail recursive optimized. *)
     let parse_targets_tr targets =
       let rec parse_targets acc = function
         | [] -> acc
-        | target :: remaining ->
-          parse_targets (uri_to_string (parse_target target) :: acc) remaining
+        | target :: remaining -> parse_targets (parse_target target :: acc) remaining
       in
       parse_targets [] targets
     ;;
+
+    let combine_installables_tr installables =
+      let rec combine_installables acc = function
+        | [] -> Some acc
+        | Nixpkgs { installables; _ } :: remaining ->
+          combine_installables (acc @ installables) remaining
+        | _ -> None
+      in
+      combine_installables [] installables
+    ;;
   end
 end

test/dune

@@ -2,7 +2,12 @@
 
 (test
  (name test_ns)
- (deps test_ns src/flake/flake.nix src/shell/shell.nix)
+ (deps
+  test_ns
+  src/flake/flake.nix
+  src/shell/shell.nix
+  src/flake_shell/flake.nix
+  src/flake_shell/shell.nix)
  (flags (:standard)))
 
 ;; Rule to compile the Rust test file into a binary

test/src/flake_shell/flake.nix

@@ -0,0 +1,12 @@
+{
+  inputs.nixpkgs.url = "github:nixos/nixpkgs?ref=nixos-unstable";
+  outputs = { self, nixpkgs }:
+    let
+      inherit (nixpkgs.lib) genAttrs;
+      inherit (nixpkgs.lib.systems) flakeExposed;
+      eachSystem = f: genAttrs flakeExposed (system: f (import nixpkgs { inherit system; }));
+    in
+    {
+      packages = eachSystem (pkgs: { default = pkgs.hello; });
+    };
+}

test/src/flake_shell/shell.nix

@@ -0,0 +1,6 @@
+{ pkgs ? import <nixpkgs> { } }:
+pkgs.mkShell {
+  buildInputs = builtins.attrValues {
+    inherit (pkgs) hello;
+  };
+}