Add set method to PropertyAssignment and ShorthandPropertyAssignment

src/compiler/expression/object/PropertyAssignment.ts

@@ -6,6 +6,7 @@ import { InitializerGetExpressionableNode, PropertyNamedNode, QuestionTokenableN
 import { Node } from "../../common";
 import { ShorthandPropertyAssignment } from "./ShorthandPropertyAssignment";
 import { PropertyAssignmentStructure, PropertyAssignmentSpecificStructure } from "../../../structures";
+import { callBaseSet } from "../../callBaseSet";
 import { callBaseGetStructure } from "../../callBaseGetStructure";
 
 // This node only has a question token in order to tell the user about bad code.
@@ -64,6 +65,21 @@ export class PropertyAssignment extends PropertyAssignmentBase<ts.PropertyAssign
         removeCommaSeparatedChild(this);
     }
 
+    /**
+     * Sets the node from a structure.
+     * @param structure - Structure to set the node with.
+     */
+    set(structure: Partial<PropertyAssignmentStructure>) {
+        callBaseSet(PropertyAssignmentBase.prototype, this, structure);
+
+        if (structure.initializer != null)
+            this.setInitializer(structure.initializer);
+        else if (structure.hasOwnProperty(nameof(structure.initializer)))
+            return this.removeInitializer();
+
+        return this;
+    }
+
     /**
      * Gets the structure equivalent to this node.
      */

src/compiler/expression/object/ShorthandPropertyAssignment.ts

@@ -7,6 +7,7 @@ import { Expression } from "../Expression";
 import { PropertyAssignment } from "./PropertyAssignment";
 import { ShorthandPropertyAssignmentStructure, ShorthandPropertyAssignmentSpecificStructure, QuestionTokenableNodeStructure } from "../../../structures";
 import { callBaseGetStructure } from "../../callBaseGetStructure";
+import { callBaseSet } from "../../callBaseSet";
 
 // This node only has an object assignment initializer, equals token, and question token, in order to tell the user about bad code
 // (See https://github.com/Microsoft/TypeScript/pull/5121/files)
@@ -97,6 +98,16 @@ export class ShorthandPropertyAssignment extends ShorthandPropertyAssignmentBase
         removeCommaSeparatedChild(this);
     }
 
+    /**
+     * Sets the node from a structure.
+     * @param structure - Structure to set the node with.
+     */
+    set(structure: Partial<ShorthandPropertyAssignmentStructure>) {
+        callBaseSet(ShorthandPropertyAssignmentBase.prototype, this, structure);
+
+        return this;
+    }
+
     /**
      * Gets the structure equivalent to this node.
      */

src/tests/compiler/expression/object/propertyAssignmentTests.ts

@@ -71,6 +71,31 @@ describe(nameof(PropertyAssignment), () => {
         });
     });
 
+    describe(nameof<PropertyAssignment>(p => p.set), () => {
+        function test(code: string, structure: Partial<PropertyAssignmentStructure>, expectedText: string) {
+            const { descendant, sourceFile } = getInfoFromTextWithDescendant<PropertyAssignment>(code, SyntaxKind.PropertyAssignment);
+            expect(descendant.set(structure).wasForgotten()).to.be.false;
+            expect(sourceFile.getFullText()).to.equal(expectedText);
+        }
+
+        it("should not change when nothing is specified", () => {
+            const code = "const t = { prop1: 1 };";
+            test(code, { }, code);
+        });
+
+        it("should set when everything is specified", () => {
+            const structure: MakeRequired<PropertyAssignmentStructure> = {
+                name: "NewName",
+                initializer: "5"
+            };
+            test("const t = { prop1: 1 };", structure, "const t = { NewName: 5 };");
+        });
+
+        it("should remove initializer when specifying undefined", () => {
+            test("const t = { prop: 1 };", { initializer: undefined }, "const t = { prop };");
+        });
+    });
+
     describe(nameof<PropertyAssignment>(p => p.getStructure), () => {
         function test(code: string, expectedStructure: MakeRequired<PropertyAssignmentStructure>) {
             const { descendant } = getInfoFromTextWithDescendant<PropertyAssignment>(code, SyntaxKind.PropertyAssignment);

src/tests/compiler/expression/object/shorthandPropertyAssignmentTests.ts

@@ -143,6 +143,26 @@ describe(nameof(ShorthandPropertyAssignment), () => {
         });
     });
 
+    describe(nameof<ShorthandPropertyAssignment>(p => p.set), () => {
+        function test(code: string, structure: Partial<ShorthandPropertyAssignmentStructure>, expectedText: string) {
+            const { descendant, sourceFile } = getInfoFromTextWithDescendant<ShorthandPropertyAssignment>(code, SyntaxKind.ShorthandPropertyAssignment);
+            expect(descendant.set(structure).wasForgotten()).to.be.false;
+            expect(sourceFile.getFullText()).to.equal(expectedText);
+        }
+
+        it("should not change when nothing is specified", () => {
+            const code = "const t = { prop1 };";
+            test(code, {}, code);
+        });
+
+        it("should set when everything is specified", () => {
+            const structure: MakeRequired<ShorthandPropertyAssignmentStructure> = {
+                name: "NewName"
+            };
+            test("const t = { prop1 };", structure, "const t = { NewName };");
+        });
+    });
+
     describe(nameof<ShorthandPropertyAssignment>(p => p.getStructure), () => {
         function test(code: string, expectedStructure: MakeRequired<ShorthandPropertyAssignmentStructure>) {
             const { descendant } = getInfoFromTextWithDescendant<ShorthandPropertyAssignment>(code, SyntaxKind.ShorthandPropertyAssignment);