Add PickProperties type and @withPickedProperties decorator (#3488)

Add PickProperties based on OmitProperties Based on #3484 (comment) you may not be wanting more mutating decorators, but PR is here if it's ok. --------- Co-authored-by: Timothee Guerin <timothee.guerin@outlook.com>
microsoft · Jun 11, 2024 · d33090e · d33090e
1 parent 22372f9
commit d33090e
Show file tree

Hide file tree

Showing 9 changed files with 137 additions and 0 deletions.
diff --git a/.chronus/changes/pick-properties-2024-5-3-17-55-34.md b/.chronus/changes/pick-properties-2024-5-3-17-55-34.md
@@ -0,0 +1,8 @@
+---
+# Change versionKind to one of: internal, fix, dependencies, feature, deprecation, breaking
+changeKind: fix
+packages:
+  - "@typespec/compiler"
+---
+
+Add `PickProperties` type to dynamically select a subset of a model 
diff --git a/docs/standard-library/built-in-data-types.md b/docs/standard-library/built-in-data-types.md
@@ -96,6 +96,23 @@ model OptionalProperties<Source>
 | Source | An object whose spread properties are all optional. |
 
 
+#### Properties
+None
+
+### `PickProperties` {#PickProperties}
+
+Represents a collection of properties with only the specified keys included.
+```typespec
+model PickProperties<Source, Keys>
+```
+
+#### Template Parameters
+| Name | Description |
+|------|-------------|
+| Source | An object whose properties are spread. |
+| Keys | The property keys to include. |
+
+
 #### Properties
 None
 

diff --git a/docs/standard-library/built-in-decorators.md b/docs/standard-library/built-in-decorators.md
@@ -933,6 +933,24 @@ Returns the model with the given properties omitted.
 
 
 
+### `@withPickedProperties` {#@withPickedProperties}
+
+Returns the model with only the given properties included.
+```typespec
+@withPickedProperties(pick: string | Union)
+```
+
+#### Target
+
+`Model`
+
+#### Parameters
+| Name | Type | Description |
+|------|------|-------------|
+| pick | `string \| Union` | List of properties to include |
+
+
+
 ### `@withUpdateableProperties` {#@withUpdateableProperties}
 
 Returns the model with non-updateable properties removed.

diff --git a/packages/compiler/generated-defs/TypeSpec.ts b/packages/compiler/generated-defs/TypeSpec.ts
@@ -76,6 +76,17 @@ export type WithoutOmittedPropertiesDecorator = (
   omit: Type
 ) => void;
 
+/**
+ * Returns the model with only the given properties included.
+ *
+ * @param pick List of properties to include
+ */
+export type WithPickedPropertiesDecorator = (
+  context: DecoratorContext,
+  target: Model,
+  pick: Type
+) => void;
+
 /**
  * Returns the model with any default values removed.
  */

diff --git a/packages/compiler/generated-defs/TypeSpec.ts-test.ts b/packages/compiler/generated-defs/TypeSpec.ts-test.ts
@@ -35,6 +35,7 @@ import {
   $visibility,
   $withDefaultKeyVisibility,
   $withOptionalProperties,
+  $withPickedProperties,
   $withUpdateableProperties,
   $withVisibility,
   $withoutDefaultValues,
@@ -76,6 +77,7 @@ import type {
   VisibilityDecorator,
   WithDefaultKeyVisibilityDecorator,
   WithOptionalPropertiesDecorator,
+  WithPickedPropertiesDecorator,
   WithUpdateablePropertiesDecorator,
   WithVisibilityDecorator,
   WithoutDefaultValuesDecorator,
@@ -88,6 +90,7 @@ type Decorators = {
   $withOptionalProperties: WithOptionalPropertiesDecorator;
   $withUpdateableProperties: WithUpdateablePropertiesDecorator;
   $withoutOmittedProperties: WithoutOmittedPropertiesDecorator;
+  $withPickedProperties: WithPickedPropertiesDecorator;
   $withoutDefaultValues: WithoutDefaultValuesDecorator;
   $withDefaultKeyVisibility: WithDefaultKeyVisibilityDecorator;
   $summary: SummaryDecorator;
@@ -131,6 +134,7 @@ const _: Decorators = {
   $withOptionalProperties,
   $withUpdateableProperties,
   $withoutOmittedProperties,
+  $withPickedProperties,
   $withoutDefaultValues,
   $withDefaultKeyVisibility,
   $summary,

diff --git a/packages/compiler/lib/std/decorators.tsp b/packages/compiler/lib/std/decorators.tsp
@@ -604,6 +604,12 @@ extern dec withoutDefaultValues(target: Model);
  */
 extern dec withoutOmittedProperties(target: Model, omit: string | Union);
 
+/**
+ * Returns the model with only the given properties included.
+ * @param pick List of properties to include
+ */
+extern dec withPickedProperties(target: Model, pick: string | Union);
+
 //---------------------------------------------------------------------------
 // Debugging
 //---------------------------------------------------------------------------

diff --git a/packages/compiler/lib/std/types.tsp b/packages/compiler/lib/std/types.tsp
@@ -53,6 +53,18 @@ model OmitProperties<Source, Keys extends string> {
   ...Source;
 }
 
+/**
+ * Represents a collection of properties with only the specified keys included.
+ *
+ * @template Source An object whose properties are spread.
+ * @template Keys The property keys to include.
+ */
+@doc("The template for picking properties.")
+@withPickedProperties(Keys)
+model PickProperties<Source, Keys extends string> {
+  ...Source;
+}
+
 /**
  * Represents a collection of properties with default values omitted.
  *

diff --git a/packages/compiler/src/lib/decorators.ts b/packages/compiler/src/lib/decorators.ts
@@ -31,6 +31,7 @@ import type {
   VisibilityDecorator,
   WithDefaultKeyVisibilityDecorator,
   WithOptionalPropertiesDecorator,
+  WithPickedPropertiesDecorator,
   WithUpdateablePropertiesDecorator,
   WithVisibilityDecorator,
   WithoutDefaultValuesDecorator,
@@ -879,6 +880,29 @@ export const $withoutOmittedProperties: WithoutOmittedPropertiesDecorator = (
   filterModelPropertiesInPlace(target, (prop) => !omitNames.has(prop.name));
 };
 
+// -- @withPickedProperties decorator ----------------------
+
+export const $withPickedProperties: WithPickedPropertiesDecorator = (
+  context: DecoratorContext,
+  target: Model,
+  pickedProperties: Type
+) => {
+  // Get the property or properties to pick
+  const pickedNames = new Set<string>();
+  if (pickedProperties.kind === "String") {
+    pickedNames.add(pickedProperties.value);
+  } else if (pickedProperties.kind === "Union") {
+    for (const variant of pickedProperties.variants.values()) {
+      if (variant.type.kind === "String") {
+        pickedNames.add(variant.type.value);
+      }
+    }
+  }
+
+  // Remove all properties not picked
+  filterModelPropertiesInPlace(target, (prop) => pickedNames.has(prop.name));
+};
+
 // -- @withoutDefaultValues decorator ----------------------
 
 export const $withoutDefaultValues: WithoutDefaultValuesDecorator = (

diff --git a/packages/compiler/test/decorators/decorators.test.ts b/packages/compiler/test/decorators/decorators.test.ts
@@ -782,6 +782,43 @@ describe("compiler: built-in decorators", () => {
     });
   });
 
+  describe("@withPickedProperties", () => {
+    it("picks a model property when given a string literal", async () => {
+      const { TestModel } = await runner.compile(
+        `
+        model OriginalModel {
+          pickMe: string;
+          notMe: string;
+        }
+
+        @test
+        model TestModel is PickProperties<OriginalModel, "pickMe"> {
+        }`
+      );
+
+      const properties = TestModel.kind === "Model" ? Array.from(TestModel.properties.keys()) : [];
+      deepStrictEqual(properties, ["pickMe"]);
+    });
+
+    it("picks model properties when given a union containing strings", async () => {
+      const { TestModel } = await runner.compile(
+        `
+        model OriginalModel {
+          pickMe: string;
+          pickMeToo: string;
+          notMe: string;
+        }
+
+        @test
+        model TestModel is PickProperties<OriginalModel, "pickMe" | "pickMeToo"> {
+        }`
+      );
+
+      const properties = TestModel.kind === "Model" ? Array.from(TestModel.properties.keys()) : [];
+      deepStrictEqual(properties, ["pickMe", "pickMeToo"]);
+    });
+  });
+
   describe("@withDefaultKeyVisibility", () => {
     it("sets the default visibility on a key property when not already present", async () => {
       const { TestModel } = (await runner.compile(