feat(ts): add flattenSubclassesToUnion param (#78)

Signed-off-by: Matt Roberts <[email protected]>
accordproject · Nov 24, 2023 · af903c3 · af903c3
1 parent e38396e
commit af903c3
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/lib/codegen/fromcto/typescript/typescriptvisitor.js b/lib/codegen/fromcto/typescript/typescriptvisitor.js
@@ -234,7 +234,6 @@ class TypescriptVisitor {
         if (field.isOptional()) {
             optional = '?';
         }
-
         const isEnumRef = field.isPrimitive() ? false
             : field.getParent().getModelFile().getModelManager().getType(field.getFullyQualifiedTypeName()).isEnum();
 
@@ -249,7 +248,16 @@ class TypescriptVisitor {
             decoratorArguments.length > 0 && (literal = ` = ${field.getType()}.${decoratorArguments}`);
         }
 
-        const tsType = this.toTsType(field.getType(), !isEnumRef && !hasUnion && !isMapRef, hasUnion);
+        let tsType = this.toTsType(field.getType(), !isEnumRef && !hasUnion && !isMapRef, hasUnion);
+
+        // If there exists direct subclasses for this field's declaration then use the union type instead
+        if (!!parameters.flattenSubclassesToUnion & !field.isPrimitive()) {
+            const subclasses = field.getParent().getModelFile().getModelManager().getType(field.getFullyQualifiedTypeName()).getDirectSubclasses();
+            if (subclasses && subclasses.length > 0) {
+                const useUnion = !(isEnumRef || isMapRef);
+                tsType = this.toTsType(field.getType(), !useUnion, useUnion);
+            }
+        }
 
         if (literal) {
             parameters.fileWriter.writeLine(1, field.getName() + literal + ';');

diff --git a/test/codegen/fromcto/typescript/typescriptvisitor.js b/test/codegen/fromcto/typescript/typescriptvisitor.js
@@ -624,6 +624,28 @@ describe('TypescriptVisitor', function () {
 
             param.fileWriter.writeLine.withArgs(1, 'literalTest = EnumType.MyEnumValue;').calledOnce.should.be.ok;
         });
+
+        it('should write a line for field name using a union type when the flattenSubclassesToUnion parameter is set', () => {
+            const mockField = sinon.createStubInstance(Field);
+            mockField.isPrimitive.returns(false);
+            mockField.getName.returns('flattenSubclassesTest');
+            mockField.getType.returns('Animal');
+            mockField.getDecorators.returns([]);
+
+            const mockModelManager = sinon.createStubInstance(ModelManager);
+            const mockModelFile = sinon.createStubInstance(ModelFile);
+            const mockClassDeclaration = sinon.createStubInstance(ClassDeclaration);
+            mockClassDeclaration.getDirectSubclasses.returns(['blah']); // Not valid, but sufficient for this test
+
+            mockModelManager.getType.returns(mockClassDeclaration);
+            mockClassDeclaration.isEnum.returns(false);
+            mockModelFile.getModelManager.returns(mockModelManager);
+            mockClassDeclaration.getModelFile.returns(mockModelFile);
+            mockField.getParent.returns(mockClassDeclaration);
+            typescriptVisitor.visitField(mockField, { ...param, flattenSubclassesToUnion: true });
+
+            param.fileWriter.writeLine.withArgs(1, 'flattenSubclassesTest: AnimalUnion;').calledOnce.should.be.ok;
+        });
     });
 
     describe('visitEnumValueDeclaration', () => {