Allow mutating namespaces

microsoft · Oct 25, 2024 · f543881 · f543881
1 parent 74ed99a
commit f543881
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git a/packages/compiler/src/experimental/mutators.ts b/packages/compiler/src/experimental/mutators.ts
@@ -73,6 +73,7 @@ export interface Mutator {
   ScalarConstructor?: MutatorRecord<ScalarConstructor>;
   StringTemplate?: MutatorRecord<StringTemplate>;
   StringTemplateSpan?: MutatorRecord<StringTemplateSpan>;
+  Namespace?: MutatorRecord<Namespace>;
 }
 
 /** @experimental */
@@ -86,7 +87,6 @@ export enum MutatorFlow {
 export type MutableType = Exclude<
   Type,
   | TemplateParameter
-  | Namespace
   | IntrinsicType
   | FunctionType
   | Decorator

diff --git a/packages/compiler/src/experimental/typekit/kits/type.ts b/packages/compiler/src/experimental/typekit/kits/type.ts
@@ -1,4 +1,4 @@
-import type { Type } from "../../../core/types.js";
+import type { Enum, Model, Type } from "../../../core/types.js";
 import { defineKit } from "../define-kit.js";
 import { copyMap } from "../utils.js";
 
@@ -63,10 +63,26 @@ defineKit<BaseTypeKit>({
         case "Enum":
           clone = this.program.checker.createType({
             ...type,
-            decorators: [...type.decorators],
             members: copyMap(type.members),
           });
           break;
+        case "Namespace":
+          clone = this.program.checker.createType({
+            ...type,
+            decorators: [...type.decorators],
+            decoratorDeclarations: new Map(type.decoratorDeclarations),
+            models:  new Map<string, Model>(type.models),
+            enums: new Map<string, Enum>(type.enums),
+            functionDeclarations: new Map(type.functionDeclarations),
+            instantiationParameters: type.instantiationParameters ? [...type.instantiationParameters] : undefined,
+            interfaces: new Map(type.interfaces),
+            namespaces: new Map(type.namespaces),
+            operations: new Map(type.operations),
+            projections: [...type.projections],
+            scalars: new Map(type.scalars),
+            unions: new Map(type.unions),
+          }) 
+        break;
         default:
           clone = this.program.checker.createType({
             ...type,

diff --git a/packages/compiler/test/experimental/mutator.test.ts b/packages/compiler/test/experimental/mutator.test.ts
@@ -1,6 +1,6 @@
 import { beforeEach, expect, it } from "vitest";
 import { mutateSubgraph, Mutator, MutatorFlow } from "../../src/experimental/mutators.js";
-import { Model } from "../../src/index.js";
+import { Model, Namespace } from "../../src/index.js";
 import { createTestHost } from "../../src/testing/test-host.js";
 import { createTestWrapper } from "../../src/testing/test-utils.js";
 import { BasicTestRunner, TestHost } from "../../src/testing/types.js";
@@ -71,6 +71,39 @@ it("recurses the model", async () => {
   expect(visited).toStrictEqual(["Foo", "Bar"]);
 });
 
+it("removes model reference from namespace", async () => {
+  const code = `
+  @test namespace Foo;
+  @test model Bar {
+    bar: string;
+  }
+  @test model Baz {
+    x: string;
+    y: string;
+    z: Bar;
+  };
+  `;
+
+const { Foo } = (await runner.compile(code)) as { Foo: Namespace, Bar: Model, Baz: Model };
+const mutator: Mutator = {
+  name: "test",
+  Namespace: {
+    mutate: (ns, clone, p, realm) => {
+      clone.models.delete("Bar");
+    }
+  },
+};
+
+const {type} = mutateSubgraph(runner.program, [mutator], Foo);
+
+const mutatedNs = type as Namespace;
+
+//Original namespace should have Bar model
+expect(Foo.models.has("Bar")).toBeTruthy();
+// Mutated namespace should not have Bar model
+expect(mutatedNs.models.has( "Bar")).toBeFalsy();
+});
+
 it("do not recurse the model", async () => {
   const code = `
     @test model Bar {