From af903c3025bad992c4f03bf6f0408bf17400a37e Mon Sep 17 00:00:00 2001
From: Matt Roberts <7544022+mttrbrts@users.noreply.github.com>
Date: Fri, 24 Nov 2023 13:04:23 +0000
Subject: [PATCH] feat(ts): add flattenSubclassesToUnion param (#78)

Signed-off-by: Matt Roberts <code@rbrts.uk>
---
 .../fromcto/typescript/typescriptvisitor.js   | 12 ++++++++--
 .../fromcto/typescript/typescriptvisitor.js   | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/lib/codegen/fromcto/typescript/typescriptvisitor.js b/lib/codegen/fromcto/typescript/typescriptvisitor.js
index 2e7d516b..9adfccf1 100644
--- 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
index 3d772097..34b7293d 100644
--- 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', () => {