feat: add Array.is type guard (#69)

manzt · Aug 10, 2023 · 5d76454 · 5d76454
1 parent 58fdbbb
commit 5d76454
Show file tree

Hide file tree

Showing 6 changed files with 254 additions and 119 deletions.
diff --git a/README.md b/README.md
@@ -41,9 +41,9 @@ import { FetchStore } from "@zarrita/storage";
 const store = new FetchStore("http://localhost:8080/data.zarr");
 
 // open array from root (note that dtype is unknown)
-const arr = await zarr.open.v2(store, { kind: "array" }); // zarr.Array<DataType, FetchStore, "/">
+const arr = await zarr.open.v2(store, { kind: "array" }); // zarr.Array<DataType, FetchStore>
 
-arr; // zarr.Array<DataType, FetchStore, "/">
+arr; // zarr.Array<DataType, FetchStore>
 arr.shape; // [5, 10]
 arr.chunk_shape; // [2, 5]
 arr.dtype; // "int32"

diff --git a/packages/core/__tests__/hierarchy.test.ts b/packages/core/__tests__/hierarchy.test.ts
@@ -0,0 +1,114 @@
+import { describe, expect, expectTypeOf, test } from "vitest";
+import { Array, Group } from "../src/hierarchy.js";
+import type {
+	ArrayMetadata,
+	BigintDataType,
+	Bool,
+	DataType,
+	Int8,
+	NumberDataType,
+	Raw,
+} from "../src/metadata.js";
+
+const array_metadata = {
+	zarr_format: 3,
+	node_type: "array",
+	data_type: "int8",
+	shape: [10, 10],
+	chunk_grid: {
+		name: "regular",
+		configuration: {
+			chunk_shape: [5, 5],
+		},
+	},
+	chunk_key_encoding: {
+		name: "default",
+		configuration: {
+			separator: "/",
+		},
+	},
+	codecs: [],
+	fill_value: 0,
+	attributes: { answer: 42 },
+} satisfies ArrayMetadata;
+
+describe("Array", () => {
+	test("constructor", async () => {
+		let arr = new Array(new Map(), "/", array_metadata);
+		expect({
+			shape: arr.shape,
+			chunk_shape: arr.chunk_shape,
+			dtype: arr.dtype,
+			fill_value: arr.fill_value,
+			attrs: await arr.attrs(),
+			path: arr.path,
+			codec_pipeline: arr.codec_pipeline,
+			store: arr.store,
+		}).toMatchInlineSnapshot(`
+			{
+			  "attrs": {
+			    "answer": 42,
+			  },
+			  "chunk_shape": [
+			    5,
+			    5,
+			  ],
+			  "codec_pipeline": {
+			    "decode": [Function],
+			    "encode": [Function],
+			  },
+			  "dtype": "int8",
+			  "fill_value": 0,
+			  "path": "/",
+			  "shape": [
+			    10,
+			    10,
+			  ],
+			  "store": Map {},
+			}
+		`);
+	});
+
+	test("Array.is", () => {
+		let arr = new Array(new Map(), "/", array_metadata as any);
+		expectTypeOf(arr.dtype).toMatchTypeOf<DataType>();
+		if (arr.is("bigint")) {
+			expectTypeOf(arr.dtype).toMatchTypeOf<BigintDataType>();
+		}
+		if (arr.is("number")) {
+			expectTypeOf(arr.dtype).toMatchTypeOf<NumberDataType>();
+		}
+		if (arr.is("bool")) {
+			expectTypeOf(arr.dtype).toMatchTypeOf<Bool>();
+		}
+		if (arr.is("raw")) {
+			expectTypeOf(arr.dtype).toMatchTypeOf<Raw>();
+		}
+		if (arr.is("int8")) {
+			expectTypeOf(arr.dtype).toMatchTypeOf<Int8>();
+		}
+	});
+});
+
+describe("Group", () => {
+	test("constructor", async () => {
+		let grp = new Group(new Map(), "/", {
+			zarr_format: 3,
+			node_type: "group",
+			attributes: { answer: 42 },
+		});
+		expect({
+			attrs: await grp.attrs(),
+			path: grp.path,
+			store: grp.store,
+		}).toMatchInlineSnapshot(`
+			{
+			  "attrs": {
+			    "answer": 42,
+			  },
+			  "path": "/",
+			  "store": Map {},
+			}
+		`);
+	});
+});
diff --git a/packages/core/__tests__/is-dtype.test.ts b/packages/core/__tests__/is-dtype.test.ts
diff --git a/packages/core/__tests__/util.test.ts b/packages/core/__tests__/util.test.ts
@@ -1,6 +1,11 @@
 import { describe, expect, test } from "vitest";
 import type { DataType } from "../src/index.js";
-import { byteswap_inplace, get_ctr, get_strides } from "../src/util.js";
+import {
+	byteswap_inplace,
+	get_ctr,
+	get_strides,
+	is_dtype,
+} from "../src/util.js";
 
 import {
 	BoolArray,
@@ -63,3 +68,73 @@ describe("get_strides", () => {
 		expect(get_strides(shape, order)).toStrictEqual(expected);
 	});
 });
+
+describe("is_dtype", () => {
+	test.each<[DataType, boolean]>([
+		["int8", true],
+		["int16", true],
+		["int32", true],
+		["uint8", true],
+		["uint16", true],
+		["uint32", true],
+		["float32", true],
+		["float64", true],
+		["bool", false],
+		["int64", false],
+		["uint64", false],
+		["r42", false],
+	])("is_dtype(%s, 'number') -> %s", (dtype, expected) => {
+		expect(is_dtype(dtype, "number")).toBe(expected);
+	});
+
+	test.each<[DataType, boolean]>([
+		["int8", false],
+		["int16", false],
+		["int32", false],
+		["uint8", false],
+		["uint16", false],
+		["uint32", false],
+		["float32", false],
+		["float64", false],
+		["bool", false],
+		["int64", true],
+		["uint64", true],
+		["r42", false],
+	])("is_dtype(%s, 'bigint') -> %s", (dtype, expected) => {
+		expect(is_dtype(dtype, "bigint")).toBe(expected);
+	});
+
+	test.each<[DataType, boolean]>([
+		["int8", false],
+		["int16", false],
+		["int32", false],
+		["uint8", false],
+		["uint16", false],
+		["uint32", false],
+		["float32", false],
+		["float64", false],
+		["bool", false],
+		["int64", false],
+		["uint64", false],
+		["r42", true],
+	])("is_dtype(%s, 'raw') -> %s", (dtype, expected) => {
+		expect(is_dtype(dtype, "raw")).toBe(expected);
+	});
+
+	test.each<DataType>([
+		"int8",
+		"int16",
+		"int32",
+		"uint8",
+		"uint16",
+		"uint32",
+		"float32",
+		"float64",
+		"bool",
+		"int64",
+		"uint64",
+		"r42",
+	])("is_dtype(%s, %s) -> true", (dtype) => {
+		expect(is_dtype(dtype, dtype)).toBe(true);
+	});
+});
diff --git a/packages/core/src/hierarchy.ts b/packages/core/src/hierarchy.ts
@@ -7,7 +7,14 @@ import type {
 	Scalar,
 } from "./metadata.js";
 import { create_codec_pipeline } from "./codecs.js";
-import { encode_chunk_key, json_decode_object, v2_marker } from "./util.js";
+import {
+	type DataTypeQuery,
+	encode_chunk_key,
+	is_dtype,
+	json_decode_object,
+	type NarrowDataType,
+	v2_marker,
+} from "./util.js";
 import { KeyError } from "./errors.js";
 
 export class Location<Store> {
@@ -136,4 +143,27 @@ export class Array<
 		}
 		return this.#attributes ?? {};
 	}
+
+	/**
+	 * A helper method to narrow `zarr.Array` Dtype.
+	 *
+	 * ```typescript
+	 * let arr: zarr.Array<DataType, FetchStore> = zarr.open(store, { kind: "array" });
+	 *
+	 * // Option 1: narrow by scalar type (e.g. "bool", "raw", "bigint", "number")
+	 * if (arr.is("bigint")) {
+	 *   // zarr.Array<"int64" | "uint64", FetchStore>
+	 * }
+	 *
+	 * // Option 3: exact match
+	 * if (arr.is("float32")) {
+	 *   // zarr.Array<"float32", FetchStore, "/">
+	 * }
+	 * ```
+	 */
+	is<Query extends DataTypeQuery>(
+		query: Query,
+	): this is Array<NarrowDataType<Dtype, Query>, Store> {
+		return is_dtype(this.dtype, query);
+	}
 }