Search code examples
javascripttypescripttypescript-typingsecmascript-next

Typescript exclude methods from class once invoked


I have a use case where I want to exclude methods from return type of class method once they are invoked. ie, let us assume I have a class Setup having method step1, step2 and step3.

class Setup {

    step1() {
        return this;
    }
    
    step2() { 
        return this;
    }
    
    step3() { 
        return this;
    }
}

let setup = new Setup();

My use case

  1. is once step1 is invoked it should return an instance of Setup which does not have step1 method at all, and users should only get the option to select between step2 and step3 and once step2 is invoked it should only get step3, as step1 and step2 were already invoked, so that a better DX can be provided
  2. The order of execution does not matter, ie, someone can execute step3 before they execute step1.
  3. And also, I am seeking the solution to work during runtime, ie during runtime, a step once invoked should be available for invocation itself.
let setup = new Setup();

setup
    .step1()
    .step2()
    .step1(); // This should not be possible as step 1 was already invoked

I have already tried this, but after invoking step2 it shows step1 as an option again. I am aware that this is partially due to Omit taking Setup as the type from which it should exclude the key. But, I am unable to find a way to refer the current instance and exclude the current method.

export type Omit<A extends object, K extends string> = Pick<A, Exclude<keyof A, K>>

class Setup {

    step1(): Omit<Setup, 'step1'> {
        return this;
    }
    
    step2(): Omit<Setup, 'step2'>{ 
        return this;
    }
    
    step3():Omit<Setup, 'step3'>{ 
        return this;
    }
}

let setup = new Setup();

Solution

  • You want both the TypeScript to issue a compiler warning if someone tries to call a method more than once in their TypeScript code, and you want a runtime error if someone tries to call a method more than once at runtime. These goals are more or less independent, and you'll have to spend effort doing each one separately. It would be nice if you could just write the code that enforces your constraint at runtime and the compiler could just inspect that and behave accordingly at compile time... but the compiler's just not smart enough to do that. So in what follows let's look at each part separately.


    First the type system:

    type OmitSetup<K extends string> = Omit<Setup<K>, K>;
    declare class Setup<K extends string = never> {
      step1(): OmitSetup<K | "step1">;
      step2(): OmitSetup<K | "step2">;
      step3(): OmitSetup<K | "step3">;
    }
    

    The idea is to make the Setup class generic in the string-constrained type parameter K corresponding to the union of method names that should be suppressed. The default type argument is never (K = never) because when you first create a Setup you haven't suppressed any method names.

    Also, since you have step1, step2, and step3 methods declared in Setup<K>, those methods will be present on Setup<K> no matter what K is. That's why I defined OmitSetup<K>, which gives you a view into Setup<K> without the methods, using the Omit utility type, and so every time you call a method with name N, the compiler returns OmitSetup<K | N>, adding N to the list of names to suppress.

    Let's walk though how it works at compile time:

    const s = new Setup();
    // const s: Setup<never>
    const s1 = s.step1();
    // const s1: OmitSetup<"step1">
    const s12 = s1.step2();
    // const s12: OmitSetup<"step1" | "step2">
    const s123 = s12.step3();
    // const s123: OmitSetup<"step1" | "step2" | "step3">
    

    So s is a Setup<never> with nothing suppressed; when we call step1() it returns an OmitSetup<"step1">, which does not have a known step1 property. If you call step2() on that, you get an OmitSetup<"step1" | "step2">, leaving you with something that only has a known step3 method. When you call that method, you get an OmitSetup<"step1" | "step2" | "step3">, and thus all the methods are suppressed.

    That gives you the desired behavior:

    s.step1().step2().step3(); // okay
    s.step2().step1().step3(); // okay
    s.step1().step2().step1(); // error!
    // -------------> ~~~~~
    // Property 'step1' does not exist on type 'OmitSetup<"step1" | "step2">'. 
    // Did you mean 'step3'?
    

    Then at runtime:

    class Setup {
      step1() {
        console.log("step1");
        return Object.assign(new Setup(), this, { step1: undefined });
      }
      step2() {
        console.log("step2");
        return Object.assign(new Setup(), this, { step2: undefined });
      }
      step3() {
        console.log("step3");
        return Object.assign(new Setup(), this, { step3: undefined });
      }
    }
    

    Here each method returns a new object (this lets us re-use existing values without mutating their states, so you can write s.step1() a million times, because s never changes, but you can never write s.step1().step1()). The new object copies all the properties from the current one, and also explicitly sets the property corresponding to the current method to undefined, so that nobody can call it at runtime. Let's test it out:

    const s = new Setup();
    s.step1().step2().step3(); // "step1", "step2", "step3"
    s.step2().step1().step3(); // "step2", "step1", "step3"
    s.step1().step2().step1(); // "step1", "step2", RUNTIME ERROR!
    // s.step1().step2().step1 is not a function
    

    Looks good; you can call the three methods in any order, but if you try to call the same method twice you get a runtime error.


    Finally, we can marry the types to the runtime code in a single TypeScript file like this:

    type OmitSetup<K extends string> = Omit<Setup<K>, K>;
    class Setup<K extends string = never> {
      step1(): OmitSetup<K | "step1"> {
        console.log("step1");
        return Object.assign(new Setup(), this, { step1: undefined }) as any
      }
      step2(): OmitSetup<K | "step2"> {
        console.log("step2");
        return Object.assign(new Setup(), this, { step2: undefined }) as any
      }
      step3(): OmitSetup<K | "step3"> {
        console.log("step3");
        return Object.assign(new Setup(), this, { step3: undefined }) as any
      }
    };
    

    This is mostly just annotating the method return types as well as asserting the values returned as the intentionally loose any type. You actually don't need as any here to get it to compile, but I've included it to make it obvious to the reader that the implementation and the typings are independent. The compiler can't understand that Object.assign(new Setup(), this, { step3: undefined }) is of type OmitSetup<K | "step3"> so we are telling it not to worry.


    Playground link to code