Merge pull request #2823 from quantified-uncertainty/decorated-exports

Support decorators on exported vars
quantified-uncertainty · Dec 29, 2023 · 11272dc · 11272dc · vercel · Dec 29, 2023
2 parents 397fc7b + f1841e5
commit 11272dc
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/packages/squiggle-lang/__tests__/SqProject/SqProject_test.ts b/packages/squiggle-lang/__tests__/SqProject/SqProject_test.ts
@@ -37,6 +37,24 @@ test("test exports", async () => {
   expect(await runFetchExports(project, "main")).toBe("{y: 6,t: 8}");
 });
 
+test("test decorated exports", async () => {
+  const project = SqProject.create();
+  project.setSource(
+    "main",
+    `
+    @name("X")
+    export x = 5
+
+    @name("Y")
+    @description("whatever")
+    export y = 6
+  `
+  );
+  expect(await runFetchExports(project, "main")).toBe(
+    '{x: 5, with params name: "X",y: 6, with params name: "Y", description: "whatever"}'
+  );
+});
+
 describe("removing sources", () => {
   const getCommonProject = () => {
     const project = SqProject.create();

diff --git a/packages/squiggle-lang/src/expression/compile.ts b/packages/squiggle-lang/src/expression/compile.ts
@@ -105,12 +105,18 @@ function compileToContent(
           currentContext
         );
         statements.push(statement);
-        if (
-          (astStatement.type === "LetStatement" ||
-            astStatement.type === "DefunStatement") &&
-          astStatement.exported
-        ) {
-          exports.push(astStatement.variable.value);
+        {
+          let maybeExportedStatement = astStatement;
+          while (maybeExportedStatement.type === "DecoratedStatement") {
+            maybeExportedStatement = maybeExportedStatement.statement;
+          }
+          if (
+            (maybeExportedStatement.type === "LetStatement" ||
+              maybeExportedStatement.type === "DefunStatement") &&
+            maybeExportedStatement.exported
+          ) {
+            exports.push(maybeExportedStatement.variable.value);
+          }
         }
         currentContext = newContext;
       }