More accurate types for EntityRefHierarchy

Author: wookieb

src/EntityRefHierarchy.ts

@@ -1,6 +1,6 @@
 import {AnyRef} from "./AnyRef";
 
-export class EntityRefHierarchy<T extends Array<AnyRef | undefined>> {
+export class EntityRefHierarchy<T extends ReadonlyArray<AnyRef | undefined>> {
 	constructor(readonly list: T) {
 		if (this.list.filter(x => x).length < 1) {
 			throw new TypeError('EntityRefHierarchy must contain at least 1 ref');
@@ -10,7 +10,7 @@ export class EntityRefHierarchy<T extends Array<AnyRef | undefined>> {
 	}
 
 	getRef<TType extends NonNullable<T[number]>['type']>(type: TType) {
-		return this.list.find(ref => ref && ref.type === type) as (T[number] & { type: TType }) | undefined;
+		return this.list.find(ref => ref && ref.type === type) as GetRefType<T, TType>
 	}
 
 	get root() {
@@ -40,3 +40,15 @@ export class EntityRefHierarchy<T extends Array<AnyRef | undefined>> {
 	}
 }
 
+
+type GetRefType<
+	T extends readonly (AnyRef | undefined)[],
+	TType extends string
+> = {
+	[K in keyof T]:
+	// If the slot is exactly `undefined` then just keep it.
+	[T[K]] extends [undefined]
+		? T[K]
+		// Otherwise, if the non-undefined part has the right type, return the entire T[K] (which may be a union with undefined)
+		: (NonNullable<T[K]> extends { type: TType } ? T[K] : never)
+}[number];

test/EntityRefHierarchyTest.ts

@@ -1,6 +1,8 @@
 import {createFactory} from "@src/createFactory";
 import {EntityRefHierarchy} from "@src/EntityRefHierarchy";
 
+import {assert, IsExact} from "conditional-type-checks";
+
 describe('EntityRefHierarchy', () => {
 
 	const factoryWorkspace = createFactory('workspace', (id: number) => id);
@@ -22,7 +24,7 @@ describe('EntityRefHierarchy', () => {
 
 
 	describe('simple hierarchy', () => {
-		const hierarchy = new EntityRefHierarchy([REF_WORKSPACE, REF_PROJECT, REF_COLLECTION]);
+		const hierarchy = new EntityRefHierarchy([REF_WORKSPACE, REF_PROJECT, REF_COLLECTION] as const);
 
 		it('getting root and leaf refs', () => {
 			expect(hierarchy.root)
@@ -40,6 +42,17 @@ describe('EntityRefHierarchy', () => {
 				.toEqual(REF_COLLECTION);
 		});
 
+
+		it('types test', () => {
+			const workspaceRef = hierarchy.getRef('workspace');
+			const projectRef = hierarchy.getRef('project');
+			const collectionRef = hierarchy.getRef('collection');
+
+			assert<IsExact<typeof workspaceRef, ReturnType<typeof factoryWorkspace>>>(true);
+			assert<IsExact<typeof projectRef, ReturnType<typeof factoryProject>>>(true);
+			assert<IsExact<typeof collectionRef, ReturnType<typeof factoryCollection>>>(true);
+		});
+
 		it('iterating', () => {
 			expect([...hierarchy])
 				.toEqual([REF_WORKSPACE, REF_PROJECT, REF_COLLECTION]);
@@ -51,7 +64,7 @@ describe('EntityRefHierarchy', () => {
 			undefined as ReturnType<typeof factoryWorkspace> | undefined,
 			REF_PROJECT,
 			REF_COLLECTION
-		]);
+		] as const);
 
 		it('getting root and leaf refs', () => {
 			expect(hierarchy.root)
@@ -69,6 +82,16 @@ describe('EntityRefHierarchy', () => {
 				.toEqual(REF_COLLECTION);
 		});
 
+		it('types test', () => {
+			const workspaceRef = hierarchy.getRef('workspace');
+			const projectRef = hierarchy.getRef('project');
+			const collectionRef = hierarchy.getRef('collection');
+
+			assert<IsExact<typeof workspaceRef, ReturnType<typeof factoryWorkspace> | undefined>>(true);
+			assert<IsExact<typeof projectRef, ReturnType<typeof factoryProject>>>(true);
+			assert<IsExact<typeof collectionRef, ReturnType<typeof factoryCollection>>>(true);
+		});
+
 		it('iterating', () => {
 			expect([...hierarchy])
 				.toEqual([REF_PROJECT, REF_COLLECTION]);