From bf3b6d2dca301dcf164995bd00b93975cf03ce5e Mon Sep 17 00:00:00 2001
From: Bob <highpressurecarsalesman@gmail.com>
Date: Sun, 17 Oct 2021 04:44:31 -0400
Subject: [PATCH] add support for --init as a file to be loaded before
 evaluation actions (#1033)

---
 src/babashka/main.clj                         | 22 ++++++++++++++++++-
 test-resources/babashka/init_caller.clj       |  4 ++++
 test-resources/babashka/init_test.clj         |  6 +++++
 .../src_for_classpath_test/call_init_main.clj |  6 +++++
 test/babashka/main_test.clj                   | 20 +++++++++++++++++
 5 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 test-resources/babashka/init_caller.clj
 create mode 100644 test-resources/babashka/init_test.clj
 create mode 100644 test-resources/babashka/src_for_classpath_test/call_init_main.clj

diff --git a/src/babashka/main.clj b/src/babashka/main.clj
index f05eae92..1152324c 100644
--- a/src/babashka/main.clj
+++ b/src/babashka/main.clj
@@ -132,6 +132,7 @@ Global opts:
   -cp, --classpath  Classpath to use. Overrides bb.edn classpath.
   --debug           Print debug information and internal stacktrace in case of exception.
   --force           Passes -Sforce to deps.clj, forcing recalculation of the classpath.
+  --init <file>     Load file after any preloads and prior to evaluation/subcommands.
 
 Help:
 
@@ -586,6 +587,10 @@ Use bb run --help to show this help output.
           (let [options (next options)]
             (recur (next options)
                    (assoc opts-map :main (first options))))
+          ("--init")
+          (let [options (next options)]
+            (recur (next options)
+                   (assoc opts-map :init (first options))))
           ("--run")
           (parse-run-opts opts-map (next options))
           ("--tasks")
@@ -662,7 +667,7 @@ Use bb run --help to show this help output.
       (let [{version-opt :version
              :keys [:shell-in :edn-in :shell-out :edn-out
                     :help :file :command-line-args
-                    :expressions :stream?
+                    :expressions :stream? :init
                     :repl :socket-repl :nrepl
                     :debug :classpath :force?
                     :main :uberscript :describe?
@@ -762,6 +767,7 @@ Use bb run --help to show this help output.
                               (error-handler e {:expression expressions
                                                 :debug debug
                                                 :preloads preloads
+                                                :init init
                                                 :loader (:loader @cp/cp-state)}))))
             expression (str/join " " expressions) ;; this might mess with the locations...
             exit-code
@@ -775,8 +781,22 @@ Use bb run --help to show this help output.
                             (error-handler e {:expression expression
                                               :debug debug
                                               :preloads preloads
+                                              :init init
                                               :loader (:loader @cp/cp-state)})))))
                     nil))
+            exit-code
+            ;; handle --init
+            (if exit-code exit-code
+                (do (when init
+                      (try
+                        (load-file* init)
+                        (catch Throwable e
+                          (error-handler e {:expression expression
+                                            :debug debug
+                                            :preloads preloads
+                                            :init init
+                                            :loader (:loader @cp/cp-state)}))))
+                    nil))
             ;; socket REPL is start asynchronously. when no other args are
             ;; provided, a normal REPL will be started as well, which causes the
             ;; process to wait until SIGINT
diff --git a/test-resources/babashka/init_caller.clj b/test-resources/babashka/init_caller.clj
new file mode 100644
index 00000000..a04c70e8
--- /dev/null
+++ b/test-resources/babashka/init_caller.clj
@@ -0,0 +1,4 @@
+(ns init-caller
+  (:require [init-test :as i]))
+
+(i/do-a-thing)
diff --git a/test-resources/babashka/init_test.clj b/test-resources/babashka/init_test.clj
new file mode 100644
index 00000000..a19bf9e7
--- /dev/null
+++ b/test-resources/babashka/init_test.clj
@@ -0,0 +1,6 @@
+(ns init-test)
+
+(defn do-a-thing [] "foo")
+
+(defn -main [& _]
+  "Hello from init!")
diff --git a/test-resources/babashka/src_for_classpath_test/call_init_main.clj b/test-resources/babashka/src_for_classpath_test/call_init_main.clj
new file mode 100644
index 00000000..8866a81e
--- /dev/null
+++ b/test-resources/babashka/src_for_classpath_test/call_init_main.clj
@@ -0,0 +1,6 @@
+(ns call-init-main
+  (:require [init-test :as i]))
+
+(defn foobar [] (str (i/do-a-thing) "bar"))
+
+(defn -main [& _] (i/do-a-thing))
diff --git a/test/babashka/main_test.clj b/test/babashka/main_test.clj
index 1be76b81..025c013d 100644
--- a/test/babashka/main_test.clj
+++ b/test/babashka/main_test.clj
@@ -189,6 +189,26 @@
                                                 (defn bar [x y] (* x y))
                                                 (bar (foo 10 30) 3)))"))))
 
+(deftest init-test
+  (testing "init with a file"
+    (is (= "foo" (bb nil "--init" "test-resources/babashka/init_test.clj" 
+                   "-f" "test-resources/babashka/init_caller.clj"))))
+  (testing "init with eval(s)"
+    (is (= "foo" (bb nil "--init" "test-resources/babashka/init_test.clj"
+                   "-e" "(init-test/do-a-thing)"))))
+  (testing "init with main from init'ed ns"
+    (is (= "Hello from init!" (bb nil "--init" "test-resources/babashka/init_test.clj"
+                                "-m" "init-test"))))
+  (testing "init with main from another namespace"
+    (test-utils/with-config '{:paths ["test-resources/babashka/src_for_classpath_test"]}
+      (is (= "foo" (bb nil "--init" "test-resources/babashka/init_test.clj"
+                     "-m" "call-init-main")))))
+  (testing "init with a qualified function passed to --main"
+    (test-utils/with-config '{:paths ["test-resources/babashka/src_for_classpath_test"]}
+      (is (= "foobar" (bb nil "--init" "test-resources/babashka/init_test.clj"
+                        "-m" "call-init-main/foobar"))))))
+    
+
 (deftest preloads-test
   ;; THIS TEST REQUIRES:
   ;; export BABASHKA_PRELOADS='(defn __bb__foo [] "foo") (defn __bb__bar [] "bar")'