Merge pull request #10 from jupyterlite/env_vars

Support environment variables that persist beyond single commands

src/commands/builtin_command_runner.ts

@@ -21,8 +21,10 @@ export class BuiltinCommandRunner implements ICommandRunner {
       return;
     }
     const path = args[0]  // Ignore other arguments?
-    // Need to handle path of "-". Maybe previous path is in an env var?
-    context.fileSystem.FS.chdir(path)
-    // Need to set PWD env var?
+    // Need to handle path of "-". Maybe previous path is in an env var?  "OLDPWD"
+
+    const { FS } = context.fileSystem
+    FS.chdir(path)
+    context.environment.set("PWD", FS.cwd())
   }
 }

src/commands/wasm_command_runner.ts

@@ -17,20 +17,28 @@ export abstract class WasmCommandRunner implements ICommandRunner {
       noInitialRun: true,
       print: (text: string) => stdout.write(`${text}\n`),
       printErr: (text: string) => stdout.write(`${text}\n`),  // Should be stderr
+      preRun: (module: any) => {
+        // Use PROXYFS so that command sees the shared FS.
+        const FS = module.FS
+        FS.mkdir(mountpoint, 0o777)
+        FS.mount(fileSystem.PROXYFS, { root: mountpoint, fs: fileSystem.FS }, mountpoint)
+        FS.chdir(fileSystem.FS.cwd())
+
+        // Copy environment variables into command.
+        context.environment.copyIntoCommand(module.ENV)
+      },
     })
-
-    // Need to use PROXYFS so that command sees the shared FS.
-    const FS = wasm.FS
-    FS.mkdir(mountpoint, 0o777)
-    FS.mount(fileSystem.PROXYFS, { root: mountpoint, fs: fileSystem.FS }, mountpoint)
-    FS.chdir(fileSystem.FS.cwd())
     const loaded = Date.now()
 
     wasm.callMain(args)
 
+    const FS = wasm.FS
     FS.close(FS.streams[1])
     FS.close(FS.streams[2])
 
+    // Copy environment variables back from command.
+    context.environment.copyFromCommand(wasm.getEnvStrings())
+
     const end = Date.now()
     console.log(`${cmdName} load time ${loaded-start} ms, run time ${end-loaded} ms`)
   }

src/context.ts

@@ -1,3 +1,4 @@
+import { Environment } from "./environment"
 import { IFileSystem } from "./file_system"
 import { Output } from "./io"
 
@@ -9,10 +10,15 @@ export class Context {
     readonly args: string[],
     readonly fileSystem: IFileSystem,
     readonly mountpoint: string,
+    environment: Environment,
     readonly stdout: Output,
-  ) {}
+  ) {
+    this.environment = environment
+  }
 
   async flush(): Promise<void> {
     await this.stdout.flush()
   }
+
+  environment: Environment
 }

src/environment.ts

@@ -0,0 +1,48 @@
+/**
+ * Collection of environment variables that are known to a shell and are passed in and out of
+ * commands.
+ */
+export class Environment {
+  constructor() {
+    this._env.set("PS1", "\x1b[1;31mjs-shell:$\x1b[1;0m ")  // red color
+  }
+
+  /**
+   * Copy environment variables back from a command after it has run.
+   */
+  copyFromCommand(source: string[]) {
+    for (const str of source) {
+      const split = str.split("=")
+      const key = split.shift()
+      if (key && !this._ignore.has(key)) {
+        this._env.set(key, split.join("="))
+      }
+    }
+  }
+
+  /**
+   * Copy environment variables into a command before it is run.
+   */
+  copyIntoCommand(target: { [key: string]: string }) {
+    for (const [key, value] of this._env.entries()) {
+      target[key] = value
+    }
+  }
+
+  get(key: string): string | null {
+    return this._env.get(key) ?? null
+  }
+
+  getPrompt(): string {
+    return this._env.get("PS1") ?? "$ "
+  }
+
+  set(key: string, value: string) {
+    this._env.set(key, value)
+  }
+
+  private _env: Map<string, string> = new Map()
+
+  // Keys to ignore when copying back from a command's env vars.
+  private _ignore: Set<string> = new Set(["USER", "LOGNAME", "HOME", "LANG", "_"])
+}

src/shell.ts

@@ -1,5 +1,6 @@
 import { CommandRegistry } from "./command_registry"
 import { Context } from "./context"
+import { Environment } from "./environment"
 import { IFileSystem } from "./file_system"
 import { FileOutput, Output, TerminalOutput } from "./io"
 import { IOutputCallback } from "./output_callback"
@@ -14,7 +15,11 @@ export class Shell {
     this._outputCallback = outputCallback
     this._mountpoint = mountpoint;
     this._currentLine = ""
-    this._prompt = "\x1b[1;31mjs-shell:$\x1b[1;0m "  // red color
+    this._environment = new Environment()
+  }
+
+  get environment(): Environment {
+    return this._environment
   }
 
   async input(char: string): Promise<void> {
@@ -26,7 +31,7 @@ export class Shell {
       const cmdText = this._currentLine.trimStart()
       this._currentLine = ""
       await this._runCommands(cmdText)
-      await this.output(this._prompt)
+      await this.output(this._environment.getPrompt())
     } else if (code == 127) {  // Backspace
       if (this._currentLine.length > 0) {
         this._currentLine = this._currentLine.slice(0, -1)
@@ -48,7 +53,7 @@ export class Shell {
       } else if (possibles.length > 1) {
         const line = possibles.join("  ")
         // Note keep leading whitespace on current line.
-        await this.output(`\r\n${line}\r\n${this._prompt}${this._currentLine}`)
+        await this.output(`\r\n${line}\r\n${this._environment.getPrompt()}${this._currentLine}`)
       }
     } else if (code == 27) {  // Escape following by 1+ more characters
       const remainder = char.slice(1)
@@ -68,6 +73,7 @@ export class Shell {
     const { FS, PATH, ERRNO_CODES, PROXYFS } = this._fsModule;
     FS.mkdir(this._mountpoint, 0o777)
     FS.chdir(this._mountpoint)
+    this._environment.set("PWD", FS.cwd())
     this._fileSystem = { FS, PATH, ERRNO_CODES, PROXYFS }
     return this._fileSystem
   }
@@ -88,7 +94,7 @@ export class Shell {
   }
 
   async start(): Promise<void> {
-    await this.output(this._prompt)
+    await this.output(this._environment.getPrompt())
   }
 
   // Keeping this public for tests.
@@ -125,7 +131,9 @@ export class Shell {
         }
 
         const args = command.suffix.map((token) => token.value)
-        const context = new Context(args, this._fileSystem!, this._mountpoint, output)
+        const context = new Context(
+          args, this._fileSystem!, this._mountpoint, this._environment, output,
+        )
         await runner.run(cmdName, context)
 
         await context.flush()
@@ -144,7 +152,7 @@ export class Shell {
 
   private readonly _outputCallback: IOutputCallback
   private _currentLine: string
-  private _prompt: string  // Should really obtain this from env
+  private _environment: Environment
 
   private _fsModule: any
   private _fileSystem?: IFileSystem