typescript - Begränsa returvärdet på en funktion baserat på argument

Betrakta följande kod

type Foo = {
  kind: "foo";
  foo: number;
};

type Bar = {
  kind: "bar";
  bar: number;
};

type FooBar = Foo | Bar;

function foobar(arg: FooBar, kind: "foo" | "bar"): FooBar | undefined {
  if (arg.kind === kind) {
    return arg;
  }
  return undefined;
}

declare const a: FooBar;

if (a.kind === "foo") {
  a.foo;  // no error, type is correctly narrowed
}

foobar(a, "foo").foo; // error: Property 'foo' does not exist on type 'Bar'

Är det möjligt att skriva en typdeklaration för foobar-funktionen, så att returvärdet begränsas enligt kind-argumentet?

Svar #1

Vi kan extrahera sorterna genom att använda generiska parametrar för att modellera de typer som vi vill använda som 'variabler' över våra outputtyper.

I det här fallet använder jag ExtractKind<T> för att extrahera unionerna av olika typer av T.

Då kan vi använda RefineKind<T, K> - där T är typen med en kind och K är kinden att förfina till.

Jag har omslutit utdataresultatet med en odefinierad kontroll - om du vill kan jag ge ett exempel på hur man uttömmande kan kartlägga alla typer för att ta bort behovet av denna kontroll, men detta verkar ligga utanför omfattningen av den ursprungliga frågan.

Låt mig veta om du behöver ytterligare förklaring på dessa koncept!

Fullständigt exempel:

type Foo = {
  kind: "foo";
  foo: number;
};

type Bar = {
  kind: "bar";
  bar: number;
};

type FooBar = Foo | Bar;

type ExtractKind<T> = T extends { kind: infer K } ? K : never;
type RefineKind<T, K> = T extends { kind: K } ? T : never;

function foobar<K extends ExtractKind<FooBar>>(arg: FooBar, kind: K) {
  if (arg.kind === kind) {
    return arg as RefineKind<FooBar, K>;
  }
  return undefined;
}

const foo = foobar(a, "foo");

if (foo !== undefined) {
  foo.foo; // foo -> number
}

Typescript-lekplatslänk här

Svar #2

Jag tror att det du verkligen försöker göra här är att implementera en användardefinierad typvakt. I praktiken skriver du en boolean-funktion med returtypen arg is SomeType och det berättar för TypeScript att om funktionen är true så måste arg vara av typen SomeType. Och om den är false betyder det att arg inte kan vara av typen SomeType.

Här är hur den funktionen ser ut i ditt fall:

function isKind<T extends Kinds>(arg: FooBar, kind: T): arg is Extract<FooBar, {kind: T}> {
    return arg.kind === kind;
}

När du använder denna funktion inuti ett if-uttalande får du meningsfull information om typen:

function testUnknown( either: FooBar ) {
    if ( isKind(either, "bar") ) {
        // we now know that either is Bar
        const barVal = either.bar; // ok
        const fooVal = either.foo; // error
    }
    else {
        // we now know that either is Foo
        const barVal = either.bar; // error
        const fooVal = either.foo; // ok
    }
}

Jag personligen skulle föredra att skriva isKind som en dubbel pilfunktion så att du enklare kan anropa isKind("foo") inuti en array.map eller annan callback, men det är bara en fråga om kodstil och personlig preferens.

const isKind = <T extends Kinds>(kind: T) => 
    (arg: FooBar): arg is Extract<FooBar, {kind: T}> => {
        return arg.kind === kind;
    }

Din ursprungliga foobar-funktion kan använda din isKind-typvakt internt. Detta är liknande @mbdavis svar men du behöver inte fastställa typen av arg med "as" eftersom TypeScript redan förstår att typen har förfinats.

function foobar<T extends Kinds>(arg: FooBar, kind: T): Extract<FooBar, {kind: T}> | undefined {
  return isKind(arg, kind) ? arg : undefined;
}

Men jag misstänker att du egentligen inte ens behöver den här foobar-funktionen och kan istället bara kalla isKind direkt i din kod.

Typescript Playground-länk