Use discriminated union to type Shape of datasets and attributes

[axelboc]

Problem

With HasShape, introduced in #1951, asserting the shape of a dataset leads to slightly different types than before, for instance: Dataset<Shape, DType> & HasShape<ArrayShape> instead of Dataset<ArrayShape, DType>.

While the two types are equivalent, this subtle change broke our Value type due to how conditional types work. At the root of the issue is the fact that ScalarShape extends ArrayShape ([] extends number[]). This has been a long-standing issue, but it hadn't really been a problem until now.

Here's a a summary of the problem:

type Shape = ArrayShape | ScalarShape | null;
type ArrayShape = number[];
type ScalarShape = [];

interface HasShape<S extends Shape = Shape> {
  shape: S;
}

// I've removed the dtype part for simplicity
type Value<D extends HasShape> = D extends HasShape<infer S>
  ? S extends ScalarShape // (A)
    ? number
    : S extends ArrayShape
      ? number[]
      : never
  : never;

// PASS
type Test1 = Value<HasShape>; // `number | number[]`
type Test2 = Value<HasShape<ScalarShape>>; // `number`
type Test3 = Value<HasShape<ArrayShape>>; // `number[]`
type Test4 = Value<HasShape<Shape & ScalarShape>>; // `number`

// FAIL
type Test5 = Value<HasShape<Shape & ArrayShape>>; // `number | number[]` instead of `number[]`

This causes issues when doing, for instance:

assertArrayShape(dataset);
assertBoolType(dataset);
const value = useDatasetValue(dataset); // `boolean | boolean[]` instead of just `boolean[]`
value.length // => error

I thought I had found a workaround in #1951 by using [S] extends [ScalarShape] at line (A), but it wasn't quite right: instead of breaking the 5th test case above, it was breaking the first case (Test1 would resolve to string[] instead of string | string[]). I uncovered the problem when trying to type-check the axes and auxiliary_signals NeXus attribute, as we allow these attributes to be either string[] or string.

Solution

I tried to patch the Value conditional type in many ways but nothing worked... So I decided to fix the problem at the root by changing the Shape type.

Conditional types work best with discriminated unions, so that's what I went for. This is already what we use for the DType type and it works great.

Here is the gist of it:

export enum ShapeClass {
  Array = 'Array',
  Scalar = 'Scalar',
  Null = 'Null',
}

export type Shape = ArrayShape | ScalarShape | NullShape;

export interface ArrayShape {
  class: ShapeClass.Array;
  dims: number[];
}

export interface ScalarShape {
  class: ShapeClass.Scalar;
  dims: [];
}

export interface NullShape {
  class: ShapeClass.Null;
}

I think it's a beneficial change overall, both internally for stricter and more verbose type-checking, and externally for end users:

BEFORE	AFTER

[axelboc]

This empty dims is just for convenience when dealing with datasets/attributes with shape ScalarShape | ArrayShape — it means I can still access obj.shape.dims and check the length to decide what to do without TypeScript complaining. See for instance renderDims(entity.shape.dims) in EntityInfo.tsx.

[loichuder]

Looks good and also more consistent in containers where we used to rename shape to dims.

Congrats 👍