Merge pull request #111 from atomic-ehr/ref-type-family-in-type-schema

TypeSchema: Add typeFamily.complexTypes, remove resourceChildren

Author: ryukzak

src/api/mustache/generator/ViewModelFactory.ts

@@ -122,17 +122,15 @@ export class ViewModelFactory {
     }
 
     private _createChildrenFor(typeRef: Identifier, cache: ViewModelCache, nestedIn?: TypeSchema): TypeViewModel[] {
+        const schema = this.tsIndex.resolve(typeRef);
+        if (!schema || !("typeFamily" in schema)) return [];
         if (isComplexTypeIdentifier(typeRef)) {
-            return this.tsIndex
-                .resourceChildren(typeRef)
-                .filter(isComplexTypeIdentifier)
+            return (schema.typeFamily?.complexTypes ?? [])
                 .filter(this.filterPred)
                 .map((childRef: Identifier) => this._createFor(childRef, cache, nestedIn));
         }
         if (isResourceIdentifier(typeRef)) {
-            return this.tsIndex
-                .resourceChildren(typeRef)
-                .filter(isResourceIdentifier)
+            return (schema.typeFamily?.resources ?? [])
                 .filter(this.filterPred)
                 .map((childRef: Identifier) => this._createFor(childRef, cache, nestedIn));
         }

src/api/writer-generator/python.ts

@@ -365,8 +365,7 @@ export class Python extends Writer<PythonGeneratorOptions> {
     }
 
     private shouldImportResourceFamily(resource: RegularTypeSchema): boolean {
-        assert(this.tsIndex !== undefined);
-        return resource.identifier.kind === "resource" && this.tsIndex.resourceChildren(resource.identifier).length > 0;
+        return resource.identifier.kind === "resource" && (resource.typeFamily?.resources?.length ?? 0) > 0;
     }
 
     private generateExportsDeclaration(
@@ -448,8 +447,7 @@ export class Python extends Writer<PythonGeneratorOptions> {
     }
 
     private generateResourceTypeField(schema: RegularTypeSchema): void {
-        assert(this.tsIndex !== undefined);
-        const hasChildren = this.tsIndex.resourceChildren(schema.identifier).length > 0;
+        const hasChildren = (schema.typeFamily?.resources?.length ?? 0) > 0;
 
         if (hasChildren) {
             this.line(`${this.nameFormatFunction("resourceType")}: str = Field(`);
@@ -660,7 +658,7 @@ export class Python extends Writer<PythonGeneratorOptions> {
             );
         const families: Record<string, string[]> = {};
         for (const resource of this.tsIndex.collectResources()) {
-            const children: string[] = this.tsIndex.resourceChildren(resource.identifier).map((c) => c.name);
+            const children = (resource.typeFamily?.resources ?? []).map((c) => c.name);
             if (children.length > 0) {
                 const familyName = `${resource.identifier.name}Family`;
                 families[familyName] = children;

src/api/writer-generator/typescript/writer.ts

@@ -10,7 +10,6 @@ import {
     isNestedIdentifier,
     isPrimitiveIdentifier,
     isProfileTypeSchema,
-    isResourceIdentifier,
     isResourceTypeSchema,
     isSpecializationTypeSchema,
     type Name,
@@ -225,7 +224,11 @@ export class TypeScript extends Writer<TypeScriptOptions> {
         const typeFamilyFields: { fieldName: string; familyTypeName: string }[] = [];
         for (const [fieldName, field] of Object.entries(schema.fields ?? {})) {
             if (isChoiceDeclarationField(field) || !field.type) continue;
-            if (isResourceIdentifier(field.type) && tsIndex.resourceChildren(field.type).length > 0) {
+            const fieldTypeSchema = tsIndex.resolve(field.type);
+            if (
+                isSpecializationTypeSchema(fieldTypeSchema) &&
+                (fieldTypeSchema.typeFamily?.resources?.length ?? 0) > 0
+            ) {
                 typeFamilyFields.push({ fieldName: tsFieldName(fieldName), familyTypeName: field.type.name });
             }
         }
@@ -257,8 +260,7 @@ export class TypeScript extends Writer<TypeScriptOptions> {
         }
         this.curlyBlock(["export", "interface", name, extendsClause], () => {
             if (isResourceTypeSchema(schema)) {
-                const possibleResourceTypes = [schema.identifier];
-                possibleResourceTypes.push(...tsIndex.resourceChildren(schema.identifier));
+                const possibleResourceTypes = [schema.identifier, ...(schema.typeFamily?.resources ?? [])];
                 const openSetSuffix =
                     this.opts.openResourceTypeSet && possibleResourceTypes.length > 1 ? " | string" : "";
                 this.lineSM(

src/typeschema/core/transformer.ts

@@ -115,11 +115,7 @@ export function extractDependencies(
     return concatIdentifiers(filtered);
 }
 
-function transformFhirSchemaResource(
-    register: Register,
-    fhirSchema: RichFHIRSchema,
-    logger?: CodegenLog,
-): TypeSchema[] {
+export function transformFhirSchema(register: Register, fhirSchema: RichFHIRSchema, logger?: CodegenLog): TypeSchema[] {
     const identifier = mkIdentifier(fhirSchema);
 
     let base: Identifier | undefined;
@@ -151,16 +147,9 @@ function transformFhirSchemaResource(
         description: fhirSchema.description,
         dependencies,
         extensions,
+        typeFamily: undefined, // NOTE: should be populateTypeFamily later.
     };
 
     const bindingSchemas = collectBindingSchemas(register, fhirSchema, logger);
     return [typeSchema, ...bindingSchemas];
 }
-
-export async function transformFhirSchema(
-    register: Register,
-    fhirSchema: RichFHIRSchema,
-    logger?: CodegenLog,
-): Promise<TypeSchema[]> {
-    return transformFhirSchemaResource(register, fhirSchema, logger);
-}

src/typeschema/index.ts

@@ -118,7 +118,7 @@ export const generateTypeSchemas = async (
             continue;
         }
 
-        for (const schema of await transformFhirSchema(register, fhirSchema, logger)) {
+        for (const schema of transformFhirSchema(register, fhirSchema, logger)) {
             schemasWithSources.push({
                 schema,
                 sourcePackage: pkgId,

src/typeschema/types.ts

@@ -94,8 +94,8 @@ type IdentifierBase = {
 };
 
 type PrimitiveIdentifier = { kind: "primitive-type" } & IdentifierBase;
-type ComplexTypeIdentifier = { kind: "complex-type" } & IdentifierBase;
-type ResourceIdentifier = { kind: "resource" } & IdentifierBase;
+export type ComplexTypeIdentifier = { kind: "complex-type" } & IdentifierBase;
+export type ResourceIdentifier = { kind: "resource" } & IdentifierBase;
 export type ValueSetIdentifier = { kind: "value-set" } & IdentifierBase;
 export type NestedIdentifier = { kind: "nested" } & IdentifierBase;
 export type BindingIdentifier = { kind: "binding" } & IdentifierBase;
@@ -271,6 +271,11 @@ export interface RegularTypeSchema {
     fields?: { [k: string]: Field };
     nested?: NestedType[];
     dependencies?: Identifier[];
+    /** Transitive children grouped by kind (e.g. Resource → { resources: [DomainResource, Patient, …] }) */
+    typeFamily?: {
+        resources?: ResourceIdentifier[];
+        complexTypes?: ComplexTypeIdentifier[];
+    };
 }
 
 export interface RegularField {

src/typeschema/utils.ts

@@ -12,9 +12,11 @@ import {
     type Identifier,
     isChoiceDeclarationField,
     isChoiceInstanceField,
+    isComplexTypeIdentifier,
     isComplexTypeTypeSchema,
     isLogicalTypeSchema,
     isProfileTypeSchema,
+    isResourceIdentifier,
     isResourceTypeSchema,
     isSpecializationTypeSchema,
     type PkgName,
@@ -103,74 +105,48 @@ export const sortAsDeclarationSequence = (schemas: RegularTypeSchema[]): Regular
 };
 
 ///////////////////////////////////////////////////////////
-// Type Schema Relations
-
-interface TypeRelation {
-    parent: Identifier;
-    child: Identifier;
-}
-
-const resourceRelatives = (schemas: TypeSchema[]): TypeRelation[] => {
-    const regularSchemas = schemas.filter(
-        (e) => isResourceTypeSchema(e) || isLogicalTypeSchema(e) || isComplexTypeTypeSchema(e),
-    );
-
-    const directPairs: TypeRelation[] = [];
-    const childrenByParent = new Map<string, Identifier[]>();
-
-    for (const schema of regularSchemas) {
-        if (schema.base) {
-            directPairs.push({ parent: schema.base, child: schema.identifier });
-            const parentName = schema.base.name;
-            let children = childrenByParent.get(parentName);
-            if (!children) {
-                children = [];
-                childrenByParent.set(parentName, children);
-            }
-            children.push(schema.identifier);
-        }
-    }
+// Type Family
 
-    const transitiveCache = new Map<string, Identifier[]>();
-    const getTransitiveChildren = (parentName: string): Identifier[] => {
-        const cached = transitiveCache.get(parentName);
-        if (cached) return cached;
+/** Populate `typeFamily` on specialization schemas with transitive children grouped by kind. */
+const populateTypeFamily = (schemas: TypeSchema[]): void => {
+    const directChildrenByParent: Record<string, Identifier[]> = {};
+    for (const schema of schemas) {
+        if (!isSpecializationTypeSchema(schema) || !schema.base) continue;
+        const parentUrl = schema.base.url;
+        if (!directChildrenByParent[parentUrl]) directChildrenByParent[parentUrl] = [];
+        directChildrenByParent[parentUrl].push(schema.identifier);
+    }
 
-        const directChildren = childrenByParent.get(parentName) ?? [];
-        const result: Identifier[] = [...directChildren];
-        for (const child of directChildren) {
-            result.push(...getTransitiveChildren(child.name));
+    const transitiveCache: Record<string, Identifier[]> = {};
+    const getTransitiveChildren = (parentUrl: string): Identifier[] => {
+        if (transitiveCache[parentUrl]) return transitiveCache[parentUrl];
+        const direct = directChildrenByParent[parentUrl] ?? [];
+        const result: Identifier[] = [...direct];
+        for (const child of direct) {
+            result.push(...getTransitiveChildren(child.url));
         }
-        transitiveCache.set(parentName, result);
+        transitiveCache[parentUrl] = result;
         return result;
     };
 
-    const seen = new Set<string>();
-    const allPairs: TypeRelation[] = [];
-
-    for (const pair of directPairs) {
-        const key = `${pair.parent.name}|${pair.child.name}`;
-        seen.add(key);
-        allPairs.push(pair);
-
-        for (const transitiveChild of getTransitiveChildren(pair.child.name)) {
-            const transitiveKey = `${pair.parent.name}|${transitiveChild.name}`;
-            if (!seen.has(transitiveKey)) {
-                seen.add(transitiveKey);
-                allPairs.push({ parent: pair.parent, child: transitiveChild });
-            }
-        }
+    for (const schema of schemas) {
+        if (!isSpecializationTypeSchema(schema)) continue;
+        const allChildren = getTransitiveChildren(schema.identifier.url);
+        if (allChildren.length === 0) continue;
+        const resources = allChildren.filter(isResourceIdentifier);
+        const complexTypes = allChildren.filter(isComplexTypeIdentifier);
+        const family: NonNullable<RegularTypeSchema["typeFamily"]> = {};
+        if (resources.length > 0) family.resources = resources;
+        if (complexTypes.length > 0) family.complexTypes = complexTypes;
+        if (Object.keys(family).length > 0) schema.typeFamily = family;
     }
-
-    return allPairs;
 };
 
 ///////////////////////////////////////////////////////////
 // Type Schema Index
 
 export type TypeSchemaIndex = {
     _schemaIndex: Record<CanonicalUrl, Record<PkgName, TypeSchema>>;
-    _relations: TypeRelation[];
     schemas: TypeSchema[];
     schemasByPackage: Record<PkgName, TypeSchema[]>;
     register?: Register;
@@ -180,7 +156,6 @@ export type TypeSchemaIndex = {
     collectProfiles: () => ProfileTypeSchema[];
     resolve: (id: Identifier) => TypeSchema | undefined;
     resolveByUrl: (pkgName: PkgName, url: CanonicalUrl) => TypeSchema | undefined;
-    resourceChildren: (id: Identifier) => Identifier[];
     tryHierarchy: (schema: TypeSchema) => TypeSchema[] | undefined;
     hierarchy: (schema: TypeSchema) => TypeSchema[];
     findLastSpecialization: (schema: TypeSchema) => TypeSchema;
@@ -241,7 +216,7 @@ export const mkTypeSchemaIndex = (
     for (const schema of schemas) {
         append(schema);
     }
-    const relations = resourceRelatives(schemas);
+    populateTypeFamily(schemas);
 
     const resolve = (id: Identifier) => {
         if (id.kind === "nested") return nestedIndex[id.url]?.[id.package];
@@ -269,10 +244,6 @@ export const mkTypeSchemaIndex = (
         return undefined;
     };
 
-    const resourceChildren = (id: Identifier): Identifier[] => {
-        return relations.filter((relative) => relative.parent.name === id.name).map((relative) => relative.child);
-    };
-
     const tryHierarchy = (schema: TypeSchema): TypeSchema[] | undefined => {
         const res: TypeSchema[] = [];
         let cur: TypeSchema | undefined = schema;
@@ -476,7 +447,6 @@ export const mkTypeSchemaIndex = (
 
     return {
         _schemaIndex: index,
-        _relations: relations,
         schemas,
         schemasByPackage: groupByPackages(schemas),
         register,
@@ -486,7 +456,6 @@ export const mkTypeSchemaIndex = (
         collectProfiles: () => schemas.filter(isProfileTypeSchema),
         resolve,
         resolveByUrl,
-        resourceChildren,
         tryHierarchy,
         hierarchy,
         findLastSpecialization,