From 674d3f59fd6204941f94f91f45459114e221ec4e Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Wed, 18 Oct 2023 11:56:06 -0700
Subject: [PATCH] Calculator field default can accept numbers

---
 .../__tests__/reducer/frTypes_test.ts           | 17 +++++++++++++++++
 packages/squiggle-lang/src/fr/calculator.ts     | 15 +++++++++++++--
 .../src/library/registry/frTypes.ts             |  6 ++++++
 .../website/src/pages/docs/Api/Calculator.mdx   |  2 +-
 4 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/packages/squiggle-lang/__tests__/reducer/frTypes_test.ts b/packages/squiggle-lang/__tests__/reducer/frTypes_test.ts
index 1ce1d0ec79..03b5b395bf 100644
--- a/packages/squiggle-lang/__tests__/reducer/frTypes_test.ts
+++ b/packages/squiggle-lang/__tests__/reducer/frTypes_test.ts
@@ -13,6 +13,7 @@ import {
   frDict,
   frOptional,
   frAny,
+  frNumberOrString,
 } from "../../src/library/registry/frTypes.js";
 import { ImmutableMap } from "../../src/utility/immutableMap.js";
 
@@ -80,6 +81,22 @@ describe("frDistOrNumber", () => {
   });
 });
 
+describe("frNumberOrString", () => {
+  test("number", () => {
+    const number = 123;
+    const value = vNumber(number);
+    expect(frNumberOrString.unpack(value)).toBe(number);
+    expect(frNumberOrString.pack(number)).toEqual(value);
+  });
+
+  test("string", () => {
+    const string = "foo";
+    const value = vString(string);
+    expect(frNumberOrString.unpack(value)).toBe(string);
+    expect(frNumberOrString.pack(string)).toEqual(value);
+  });
+});
+
 describe("frDist", () => {
   const dResult = Normal.make({ mean: 2, stdev: 5 });
   if (!dResult.ok) {
diff --git a/packages/squiggle-lang/src/fr/calculator.ts b/packages/squiggle-lang/src/fr/calculator.ts
index 038125a8af..a91bfca8bc 100644
--- a/packages/squiggle-lang/src/fr/calculator.ts
+++ b/packages/squiggle-lang/src/fr/calculator.ts
@@ -5,6 +5,7 @@ import {
   frArray,
   frString,
   frOptional,
+  frNumberOrString,
 } from "../library/registry/frTypes.js";
 import { FnFactory } from "../library/registry/helpers.js";
 import { vCalculator } from "../value/index.js";
@@ -14,6 +15,16 @@ const maker = new FnFactory({
   requiresNamespace: true,
 });
 
+const convertFieldDefault = (value: number | string | null): string => {
+  if (typeof value === "number") {
+    return value.toString();
+  } else if (typeof value === "string") {
+    return value;
+  } else {
+    return "";
+  }
+};
+
 export const library = [
   maker.make({
     name: "make",
@@ -30,7 +41,7 @@ export const library = [
               frArray(
                 frDict(
                   ["name", frString],
-                  ["default", frOptional(frString)],
+                  ["default", frOptional(frNumberOrString)],
                   ["description", frOptional(frString)]
                 )
               ),
@@ -43,7 +54,7 @@ export const library = [
             description: description || undefined,
             fields: fields.map((vars) => ({
               name: vars.name,
-              default: vars.default || "",
+              default: convertFieldDefault(vars.default),
               description: vars.description || undefined,
             })),
           });
diff --git a/packages/squiggle-lang/src/library/registry/frTypes.ts b/packages/squiggle-lang/src/library/registry/frTypes.ts
index b65a7a43fc..107782ebb3 100644
--- a/packages/squiggle-lang/src/library/registry/frTypes.ts
+++ b/packages/squiggle-lang/src/library/registry/frTypes.ts
@@ -57,6 +57,12 @@ export const frDistOrNumber: FRType<BaseDist | number> = {
   pack: (v) => (typeof v === "number" ? vNumber(v) : vDist(v)),
   getName: () => "distribution|number",
 };
+export const frNumberOrString: FRType<string | number> = {
+  unpack: (v) =>
+    v.type === "String" ? v.value : v.type === "Number" ? v.value : undefined,
+  pack: (v) => (typeof v === "number" ? vNumber(v) : vString(v)),
+  getName: () => "number|string",
+};
 export const frDist: FRType<BaseDist> = {
   unpack: (v) => (v.type === "Dist" ? v.value : undefined),
   pack: (v) => vDist(v),
diff --git a/packages/website/src/pages/docs/Api/Calculator.mdx b/packages/website/src/pages/docs/Api/Calculator.mdx
index 5f09d44f78..2adcf47b82 100644
--- a/packages/website/src/pages/docs/Api/Calculator.mdx
+++ b/packages/website/src/pages/docs/Api/Calculator.mdx
@@ -18,7 +18,7 @@ Calculator.make: ({
   description: string,
   fields: list<{
     name: string,
-    default?: string,
+    default?: string | number,
     description?: string
   }>
 }) => calculator