From 979e46248693aab8722a197ce5c1f1e70788a6eb Mon Sep 17 00:00:00 2001
From: Tom McLaughlin <tom@codedown.io>
Date: Fri, 1 Dec 2023 16:47:56 -0700
Subject: [PATCH] Be able to pass extra flags to stack integration

---
 main/Main.hs            | 5 ++++-
 src/IHaskell/Flags.hs   | 8 +++++++-
 src/IHaskell/IPython.hs | 3 +++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/main/Main.hs b/main/Main.hs
index 5161b4dc..2608d9db 100644
--- a/main/Main.hs
+++ b/main/Main.hs
@@ -113,6 +113,8 @@ parseKernelArgs = foldl' addFlag defaultKernelSpecOptions
       kernelSpecOpts { kernelSpecInstallPrefix = Just prefix }
     addFlag kernelSpecOpts KernelspecUseStack =
       kernelSpecOpts { kernelSpecUseStack = True }
+    addFlag kernelSpecOpts (KernelspecStackFlag flag) =
+      kernelSpecOpts { kernelSpecStackFlags = flag : (kernelSpecStackFlags kernelSpecOpts) }
     addFlag kernelSpecOpts (KernelspecEnvFile fp) =
       kernelSpecOpts { kernelSpecEnvFile = Just fp }
     addFlag _kernelSpecOpts flag = error $ "Unknown flag" ++ show flag
@@ -125,6 +127,7 @@ runKernel kOpts profileSrc = do
   let debug = kernelSpecDebug kOpts
       libdir = kernelSpecGhcLibdir kOpts
       useStack = kernelSpecUseStack kOpts
+      stackFlags = kernelSpecStackFlags kOpts
 
   -- Parse the profile file.
   let profileErr = error $ "ihaskell: "++profileSrc++": Failed to parse profile file"
@@ -145,7 +148,7 @@ runKernel kOpts profileSrc = do
     -- If we're in a stack directory, use `stack` to set the environment
     -- We can't do this with base <= 4.6 because setEnv doesn't exist.
     when stack $
-      readProcess "stack" ["exec", "env"] "" >>= parseAndSetEnv
+      readProcess "stack" (["exec", "env"] <> stackFlags) "" >>= parseAndSetEnv
 
   case kernelSpecEnvFile kOpts of
     Nothing -> return ()
diff --git a/src/IHaskell/Flags.hs b/src/IHaskell/Flags.hs
index cdd91a2e..5ce9c258 100644
--- a/src/IHaskell/Flags.hs
+++ b/src/IHaskell/Flags.hs
@@ -43,6 +43,7 @@ data Argument = ConfFile String               -- ^ A file with commands to load
               | ConvertLhsStyle (LhsStyle String)
               | KernelspecInstallPrefix String
               | KernelspecUseStack
+              | KernelspecStackFlag String
               | KernelspecEnvFile FilePath
   deriving (Eq, Show)
 
@@ -153,6 +154,10 @@ kernelStackFlag = flagNone ["stack"] addStack
   where
     addStack (Args md prev) = Args md (KernelspecUseStack : prev)
 
+kernelStackExtraFlags :: Flag Args
+kernelStackExtraFlags = flagReq ["stack-flag"] (store KernelspecStackFlag) ""
+        "Extra flag to pass to `stack` when --stack is used. Can be specified multiple times."
+
 kernelEnvFileFlag :: Flag Args
 kernelEnvFileFlag =
   flagReq
@@ -181,7 +186,7 @@ store constructor str (Args md prev) = Right $ Args md $ constructor str : prev
 installKernelSpec :: Mode Args
 installKernelSpec =
   mode "install" (Args InstallKernelSpec []) "Install the Jupyter kernelspec." noArgs
-    [ghcLibFlag, ghcRTSFlag, kernelDebugFlag, kernelNameFlag, displayNameFlag, confFlag, installPrefixFlag, helpFlag, kernelStackFlag, kernelEnvFileFlag]
+    [ghcLibFlag, ghcRTSFlag, kernelDebugFlag, kernelNameFlag, displayNameFlag, confFlag, installPrefixFlag, helpFlag, kernelStackFlag, kernelStackExtraFlags, kernelEnvFileFlag]
 
 kernel :: Mode Args
 kernel = mode "kernel" (Args (Kernel Nothing) []) "Invoke the IHaskell kernel." kernelArg
@@ -189,6 +194,7 @@ kernel = mode "kernel" (Args (Kernel Nothing) []) "Invoke the IHaskell kernel."
            , kernelDebugFlag
            , confFlag
            , kernelStackFlag
+           , kernelStackExtraFlags
            , kernelEnvFileFlag
            , kernelCodeMirrorFlag
            , kernelHtmlCodeWrapperClassFlag
diff --git a/src/IHaskell/IPython.hs b/src/IHaskell/IPython.hs
index 2fb04b13..b3b73179 100644
--- a/src/IHaskell/IPython.hs
+++ b/src/IHaskell/IPython.hs
@@ -45,6 +45,7 @@ data KernelSpecOptions =
          , kernelSpecConfFile :: IO (Maybe String)        -- ^ Filename of profile JSON file.
          , kernelSpecInstallPrefix :: Maybe String
          , kernelSpecUseStack :: Bool                     -- ^ Whether to use @stack@ environments.
+         , kernelSpecStackFlags :: [String]               -- ^ Extra flags to pass to @stack@.
          , kernelSpecEnvFile :: Maybe FilePath
          , kernelSpecKernelName :: String                 -- ^ The IPython kernel name
          , kernelSpecDisplayName :: String                -- ^ The IPython kernel display name
@@ -62,6 +63,7 @@ defaultKernelSpecOptions = KernelSpecOptions
   , kernelSpecConfFile = defaultConfFile
   , kernelSpecInstallPrefix = Nothing
   , kernelSpecUseStack = False
+  , kernelSpecStackFlags = []
   , kernelSpecEnvFile = Nothing
   , kernelSpecKernelName = "haskell"
   , kernelSpecDisplayName = "Haskell"
@@ -150,6 +152,7 @@ installKernelspec repl opts = void $ do
              [] -> []
              _ -> "+RTS" : kernelSpecRTSOptions opts ++ ["-RTS"])
            ++ ["--stack" | kernelSpecUseStack opts]
+           ++ mconcat [["--stack-flag", f] | f <- kernelSpecStackFlags opts]
 
   let kernelSpec = KernelSpec
         { kernelDisplayName = kernelSpecDisplayName opts