refactor(language-service): Create helper methods for manipulating object and array ASTs. (#47181)

dylhunn · alxhub · commit 7f74e35ffc3f · 2022-08-22T10:44:22.000-07:00
Create three new helper methods: `addElementToArrayLiteral`, `objectPropertyAssignmentForKey`, and `updateObjectValueForKey`. These methods make interacting with array and object literals easier. These will be useful for the standalone imports feature, which will need to add new terms to import arrays in Component and NgModule decorators. PR Close #47181
diff --git a/packages/language-service/src/ts_utils.ts b/packages/language-service/src/ts_utils.ts
@@ -94,3 +94,38 @@ export function collectMemberMethods(
   }
   return members;
 }
+
+/**
+ * Given an existing array literal expression, update it by pushing a new expression.
+ */
+export function addElementToArrayLiteral(
+    arr: ts.ArrayLiteralExpression, elem: ts.Expression): ts.ArrayLiteralExpression {
+  return ts.factory.updateArrayLiteralExpression(arr, [...arr.elements, elem]);
+}
+
+/**
+ * Given an ObjectLiteralExpression node, extract and return the PropertyAssignment corresponding to
+ * the given key. `null` if no such key exists.
+ */
+export function objectPropertyAssignmentForKey(
+    obj: ts.ObjectLiteralExpression, key: string): ts.PropertyAssignment|null {
+  const matchingProperty = obj.properties.filter(
+      a => a.name !== undefined && ts.isIdentifier(a.name) && a.name.escapedText === key)[0];
+  return matchingProperty && ts.isPropertyAssignment(matchingProperty) ? matchingProperty : null;
+}
+
+/**
+ * Given an ObjectLiteralExpression node, create or update the specified key, using the provided
+ * callback to generate the new value (possibly based on an old value).
+ */
+export function updateObjectValueForKey(
+    obj: ts.ObjectLiteralExpression, key: string,
+    newValueFn: (oldValue?: ts.Expression) => ts.Expression): ts.ObjectLiteralExpression {
+  const existingProp = objectPropertyAssignmentForKey(obj, key);
+  const newProp = ts.factory.createPropertyAssignment(
+      ts.factory.createIdentifier(key), newValueFn(existingProp?.initializer));
+  return ts.factory.updateObjectLiteralExpression(obj, [
+    ...obj.properties.filter(p => p !== existingProp),
+    newProp,
+  ]);
+}
diff --git a/packages/language-service/test/ts_utils_spec.ts b/packages/language-service/test/ts_utils_spec.ts
@@ -9,10 +9,10 @@
 import {initMockFileSystem} from '@angular/compiler-cli/src/ngtsc/file_system/testing';
 import ts from 'typescript';
 
-import {collectMemberMethods, findTightestNode} from '../src/ts_utils';
+import {addElementToArrayLiteral, collectMemberMethods, findTightestNode, objectPropertyAssignmentForKey, updateObjectValueForKey} from '../src/ts_utils';
 import {LanguageServiceTestEnv, OpenBuffer, Project} from '../testing';
 
-describe('ts utils', () => {
+describe('TS util', () => {
   describe('collectMemberMethods', () => {
     beforeEach(() => {
       initMockFileSystem('Native');
@@ -84,4 +84,82 @@ describe('ts utils', () => {
           .sort();
     }
   });
+
+  describe('AST method', () => {
+    let printer: ts.Printer;
+    let sourceFile: ts.SourceFile;
+
+    function print(node: ts.Node): string {
+      return printer.printNode(ts.EmitHint.Unspecified, node, sourceFile);
+    }
+
+    beforeAll(() => {
+      printer = ts.createPrinter();
+      sourceFile =
+          ts.createSourceFile('placeholder.ts', '', ts.ScriptTarget.ESNext, true, ts.ScriptKind.TS);
+    });
+
+    describe('addElementToArrayLiteral', () => {
+      it('transforms an empty array literal expression', () => {
+        const oldArr = ts.factory.createArrayLiteralExpression([], false);
+        const newArr = addElementToArrayLiteral(oldArr, ts.factory.createStringLiteral('a'));
+        expect(print(newArr)).toEqual('["a"]');
+      });
+
+      it('transforms an existing array literal expression', () => {
+        const oldArr =
+            ts.factory.createArrayLiteralExpression([ts.factory.createStringLiteral('a')], false);
+        const newArr = addElementToArrayLiteral(oldArr, ts.factory.createStringLiteral('b'));
+        expect(print(newArr)).toEqual('["a", "b"]');
+      });
+    });
+
+    describe('objectPropertyAssignmentForKey', () => {
+      let oldObj: ts.ObjectLiteralExpression;
+
+      beforeEach(() => {
+        oldObj = ts.factory.createObjectLiteralExpression(
+            [ts.factory.createPropertyAssignment(
+                ts.factory.createIdentifier('foo'), ts.factory.createStringLiteral('bar'))],
+            false);
+      });
+
+      it('returns null when no property exists', () => {
+        const prop = objectPropertyAssignmentForKey(oldObj, 'oops');
+        expect(prop).toBeNull();
+      });
+
+      it('returns the requested property assignment', () => {
+        const prop = objectPropertyAssignmentForKey(oldObj, 'foo');
+        expect(print(prop!)).toEqual('foo: "bar"');
+      });
+    });
+
+    describe('updateObjectValueForKey', () => {
+      let oldObj: ts.ObjectLiteralExpression;
+
+      const valueAppenderFn = (oldValue?: ts.Expression) => {
+        if (!oldValue) return ts.factory.createStringLiteral('baz');
+        if (!ts.isStringLiteral(oldValue)) return oldValue;
+        return ts.factory.createStringLiteral(oldValue.text + 'baz');
+      };
+
+      beforeEach(() => {
+        oldObj = ts.factory.createObjectLiteralExpression(
+            [ts.factory.createPropertyAssignment(
+                ts.factory.createIdentifier('foo'), ts.factory.createStringLiteral('bar'))],
+            false);
+      });
+
+      it('creates a non-existant property', () => {
+        const obj = updateObjectValueForKey(oldObj, 'newKey', valueAppenderFn);
+        expect(print(obj)).toBe('{ foo: "bar", newKey: "baz" }');
+      });
+
+      it('updates an existing property', () => {
+        const obj = updateObjectValueForKey(oldObj, 'foo', valueAppenderFn);
+        expect(print(obj)).toBe('{ foo: "barbaz" }');
+      });
+    });
+  });
 });
