Improve the type system

So that authors defining a language shouldn’t add `!` everywhere they refer to  `base`.

Author: DmitrySharabin

src/core/registry.ts

@@ -2,7 +2,13 @@ import { kebabToCamelCase } from '../shared/util';
 import { cloneGrammar } from '../util/extend';
 import { forEach, toArray } from '../util/iterables';
 import { extend } from '../util/language-util';
-import type { ComponentProto, Grammar, LanguageProto } from '../types';
+import type {
+	ComponentProto,
+	Grammar,
+	GrammarOptions,
+	GrammarOptionsWithBase,
+	LanguageProto,
+} from '../types';
 import type { Prism } from './prism';
 
 interface Entry {
@@ -172,18 +178,27 @@ export class Registry {
 		// We need this so that any code modifying the base grammar doesn't affect other instances
 		const baseGrammar = base && cloneGrammar(required(base.id), base.id);
 
-		let evaluatedGrammar =
-			typeof grammar === 'object'
-				? grammar
-				: grammar({
-						base: baseGrammar,
-						getLanguage: required,
-						getOptionalLanguage: id => this.getLanguage(id),
-						extend: (id, ref) => extend(required(id), id, ref),
-					});
-
-		if (base) {
-			evaluatedGrammar = extend(baseGrammar!, base.id, evaluatedGrammar);
+		let evaluatedGrammar: Grammar;
+		if (typeof grammar === 'object') {
+			// if the grammar is an object, we can use it directly
+			evaluatedGrammar = grammar;
+		}
+		else {
+			const options: GrammarOptions = {
+				getLanguage: required,
+				getOptionalLanguage: id => this.getLanguage(id),
+				extend: (id, ref) => extend(required(id), id, ref),
+				...(baseGrammar && { base: baseGrammar }),
+			};
+
+			const grammarFn = grammar as (
+				options: GrammarOptions | GrammarOptionsWithBase
+			) => Grammar;
+			evaluatedGrammar = grammarFn(options);
+		}
+
+		if (baseGrammar) {
+			evaluatedGrammar = extend(baseGrammar, base.id, evaluatedGrammar);
 		}
 
 		return (entry.evaluatedGrammar = evaluatedGrammar);

src/types.d.ts

@@ -3,22 +3,43 @@ import type { Prism } from './core/prism';
 
 export interface GrammarOptions {
 	readonly getLanguage: (id: string) => Grammar;
-	readonly base?: Grammar;
 	readonly getOptionalLanguage: (id: string) => Grammar | undefined;
 	readonly extend: (id: string, ref: GrammarTokens) => Grammar;
 }
+
+// Overload for when base is required
+export interface GrammarOptionsWithBase extends GrammarOptions {
+	readonly base: Grammar;
+}
+
 export interface ComponentProtoBase<Id extends string = string> {
 	id: Id;
 	require?: ComponentProto | readonly ComponentProto[];
 	optional?: string | readonly string[];
 	alias?: string | readonly string[];
 	effect?: (Prism: Prism & { plugins: Record<KebabToCamelCase<Id>, {}> }) => () => void;
 }
-export interface LanguageProto<Id extends string = string> extends ComponentProtoBase<Id> {
+
+// For languages that extend a base language
+export interface LanguageProtoWithBase<Id extends string = string> extends ComponentProtoBase<Id> {
+	grammar: Grammar | ((options: GrammarOptionsWithBase) => Grammar);
+	plugin?: undefined;
+	base: LanguageProto; // Required base
+}
+
+// For languages that don't extend a base language
+export interface LanguageProtoWithoutBase<Id extends string = string>
+	extends ComponentProtoBase<Id> {
 	grammar: Grammar | ((options: GrammarOptions) => Grammar);
 	plugin?: undefined;
-	base?: LanguageProto;
+	base?: never; // Explicitly no base allowed
 }
+
+// Union type that allows TypeScript to discriminate
+export type LanguageProto<Id extends string = string> =
+	| LanguageProtoWithBase<Id>
+	| LanguageProtoWithoutBase<Id>;
+
 type PluginType<Name extends string> = unknown;
 export interface PluginProto<Id extends string = string> extends ComponentProtoBase<Id> {
 	grammar?: undefined;

tests/core/registry.ts

@@ -1,5 +1,6 @@
 import { assert } from 'chai';
 import { Prism } from '../../src/core/prism';
+import type { GrammarOptions } from '../../src/types';
 
 describe('Registry', () => {
 	it('should resolve aliases', () => {
@@ -26,7 +27,7 @@ describe('Registry', () => {
 		components.add({
 			id: 'c',
 			optional: 'b',
-			grammar ({ getOptionalLanguage }) {
+			grammar ({ getOptionalLanguage }: GrammarOptions) {
 				return getOptionalLanguage('b') ?? {};
 			},
 		});