support jsImports

packages/components/src/components/SquiggleChart.tsx

@@ -8,14 +8,7 @@ import {
 } from "@quri/squiggle-lang";
 import { useSquiggle } from "../lib/hooks";
 import { SquiggleViewer } from "./SquiggleViewer";
-
-type jsImports =
-  | number
-  | string
-  | jsImports[]
-  | {
-      [k: string]: jsImports;
-    };
+import { JsImports } from "../lib/jsImports";
 
 export interface SquiggleChartProps {
   /** The input string for squiggle */
@@ -38,7 +31,7 @@ export interface SquiggleChartProps {
   width?: number;
   height?: number;
   /** JS imported parameters */
-  jsImports?: jsImports;
+  jsImports?: JsImports;
   /** Whether to show a summary of the distribution */
   showSummary?: boolean;
   /** Set the x scale to be logarithmic by deault */
@@ -61,7 +54,7 @@ export interface SquiggleChartProps {
 }
 
 const defaultOnChange = () => {};
-const defaultImports = {};
+const defaultImports: JsImports = {};
 
 export const SquiggleChart: React.FC<SquiggleChartProps> = React.memo(
   ({

packages/components/src/components/SquigglePlayground.tsx

@@ -39,6 +39,7 @@ import { ViewSettings, viewSettingsSchema } from "./ViewSettings";
 import { HeadedSection } from "./ui/HeadedSection";
 import { defaultTickFormat } from "../lib/distributionSpecBuilder";
 import { Button } from "./ui/Button";
+import { JsImports } from "../lib/jsImports";
 
 type PlaygroundProps = SquiggleChartProps & {
   /** The initial squiggle string to put in the playground */
@@ -112,8 +113,8 @@ const SamplingSettings: React.FC<{ register: UseFormRegister<FormFields> }> = ({
 );
 
 const InputVariablesSettings: React.FC<{
-  initialImports: any; // TODO - any json type
-  setImports: (imports: any) => void;
+  initialImports: JsImports;
+  setImports: (imports: JsImports) => void;
 }> = ({ initialImports, setImports }) => {
   const [importString, setImportString] = useState(() =>
     JSON.stringify(initialImports)
@@ -122,7 +123,7 @@ const InputVariablesSettings: React.FC<{
 
   const onChange = (value: string) => {
     setImportString(value);
-    let imports = {} as any;
+    let imports = {};
     try {
       imports = JSON.parse(value);
       setImportsAreValid(true);
@@ -251,7 +252,7 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
     onChange: onCodeChange,
   });
 
-  const [imports, setImports] = useState({});
+  const [imports, setImports] = useState<JsImports>({});
 
   const { register, control } = useForm({
     resolver: yupResolver(schema),
@@ -309,7 +310,7 @@ export const SquigglePlayground: FC<PlaygroundProps> = ({
           executionId={executionId}
           environment={env}
           {...vars}
-          // jsImports={imports}
+          jsImports={imports}
           enableLocalSettings={true}
         />
       </div>

packages/components/src/lib/hooks/useSquiggle.ts

@@ -1,18 +1,11 @@
 import { environment, SqProject, SqValue } from "@quri/squiggle-lang";
 import { useEffect, useMemo } from "react";
-
-export type jsImports =
-  | number
-  | string
-  | jsImports[]
-  | {
-      [k: string]: jsImports;
-    };
+import { JsImports, jsImportsToSquiggleCode } from "../jsImports";
 
 type SquiggleArgs = {
   code: string;
   executionId?: number;
-  jsImports?: jsImports;
+  jsImports?: JsImports;
   environment?: environment;
   onChange?: (expr: SqValue | undefined) => void;
 };
@@ -25,12 +18,10 @@ export const useSquiggle = (args: SquiggleArgs) => {
       if (args.environment) {
         project.setEnvironment(args.environment);
       }
-      if (args.jsImports) {
-        console.log(JSON.stringify(args.jsImports));
-        project.setSource(
-          "zzz", // due to bug in topology implementation, can be renamed later
-          "imports = " + JSON.stringify(args.jsImports)
-        );
+      if (args.jsImports && Object.keys(args.jsImports).length) {
+        const importsSource = jsImportsToSquiggleCode(args.jsImports);
+        project.setSource("imports", importsSource);
+        project.setContinues("main", ["imports"]);
       }
       project.run("main");
       const result = project.getResult("main");
@@ -38,12 +29,7 @@ export const useSquiggle = (args: SquiggleArgs) => {
       return { result, bindings };
     },
     // eslint-disable-next-line react-hooks/exhaustive-deps
-    [
-      args.code,
-      args.environment,
-      // args.jsImports,
-      args.executionId,
-    ]
+    [args.code, args.environment, args.jsImports, args.executionId]
   );
 
   const { onChange } = args;

packages/components/src/lib/jsImports.ts

@@ -0,0 +1,41 @@
+type JsImportsValue =
+  | number
+  | string
+  | JsImportsValue[]
+  | {
+      [k: string]: JsImportsValue;
+    };
+
+export type JsImports = {
+  [k: string]: JsImportsValue;
+};
+
+const quote = (arg: string) => `"${arg.replace(new RegExp('"', "g"), '\\"')}"`;
+
+const jsImportsValueToSquiggleCode = (v: JsImportsValue): string => {
+  if (typeof v === "number") {
+    return String(v);
+  } else if (typeof v === "string") {
+    return quote(v);
+  } else if (v instanceof Array) {
+    return "[" + v.map((x) => jsImportsValueToSquiggleCode(x)) + "]";
+  } else {
+    if (Object.keys(v).length) {
+      return (
+        "{" +
+        Object.entries(v)
+          .map(([k, v]) => `${k}:${jsImportsValueToSquiggleCode(v)},`)
+          .join("") +
+        "}"
+      );
+    } else {
+      return "0"; // squiggle doesn't support empty `{}`
+    }
+  }
+};
+
+export const jsImportsToSquiggleCode = (v: JsImports) => {
+  return Object.entries(v)
+    .map(([k, v]) => `$${k} = ${jsImportsValueToSquiggleCode(v)}\n`)
+    .join("");
+};