👽️ Update listFiles to HF's updated API

coyotte508 · coyotte508 · commit c7cd12c76363 · 2023-03-02T17:20:46.000+01:00
diff --git a/packages/hub/src/lib/list-files.spec.ts b/packages/hub/src/lib/list-files.spec.ts
@@ -22,12 +22,11 @@ describe("listFiles", () => {
 		assert.deepStrictEqual(files, [
 			{
 				lastCommit: {
-					author: {
-						date: "2018-11-14T23:35:08.000Z",
-					},
-					id:      "504939aa53e8ce310dba3dd2296dbe266c575de4",
-					subject: "initial commit",
+					date:  "2018-11-14T23:35:08.000Z",
+					id:    "504939aa53e8ce310dba3dd2296dbe266c575de4",
+					title: "initial commit",
 				},
+				oid:      "dc08351d4dc0732d9c8af04070ced089b201ce2f",
 				path:     ".gitattributes",
 				security: {
 					avScan: {
@@ -45,12 +44,11 @@ describe("listFiles", () => {
 			},
 			{
 				lastCommit: {
-					author: {
-						date: "2019-06-18T09:06:51.000Z",
-					},
-					id:      "bb3c1c3256d2598217df9889a14a2e811587891d",
-					subject: "Update config.json",
+					date:  "2019-06-18T09:06:51.000Z",
+					id:    "bb3c1c3256d2598217df9889a14a2e811587891d",
+					title: "Update config.json",
 				},
+				oid:      "fca794a5f07ff8f963fe8b61e3694b0fb7f955df",
 				path:     "config.json",
 				security: {
 					avScan: {
@@ -68,16 +66,16 @@ describe("listFiles", () => {
 			},
 			{
 				lastCommit: {
-					author: {
-						date: "2019-06-18T09:06:34.000Z",
-					},
-					id:      "3d2477d72b675a999d1b13ca822aaaf4908634ad",
-					subject: "Update pytorch_model.bin",
+					date:  "2019-06-18T09:06:34.000Z",
+					id:    "3d2477d72b675a999d1b13ca822aaaf4908634ad",
+					title: "Update pytorch_model.bin",
 				},
 				lfs: {
-					oid:  "097417381d6c7230bd9e3557456d726de6e83245ec8b24f529f60198a67b203a",
-					size: 440473133,
+					oid:         "097417381d6c7230bd9e3557456d726de6e83245ec8b24f529f60198a67b203a",
+					size:        440473133,
+					pointerSize: 134,
 				},
+				oid:      "ba5d19791be1dd7992e33bd61f20207b0f7f50a5",
 				path:     "pytorch_model.bin",
 				security: {
 					avScan: {
@@ -114,16 +112,16 @@ describe("listFiles", () => {
 			},
 			{
 				lastCommit: {
-					author: {
-						date: "2019-09-23T19:48:44.000Z",
-					},
-					id:      "dd4bc8b21efa05ec961e3efc4ee5e3832a3679c7",
-					subject: "Update tf_model.h5",
+					date:  "2019-09-23T19:48:44.000Z",
+					id:    "dd4bc8b21efa05ec961e3efc4ee5e3832a3679c7",
+					title: "Update tf_model.h5",
 				},
 				lfs: {
-					oid:  "a7a17d6d844b5de815ccab5f42cad6d24496db3850a2a43d8258221018ce87d2",
-					size: 536063208,
+					oid:         "a7a17d6d844b5de815ccab5f42cad6d24496db3850a2a43d8258221018ce87d2",
+					size:        536063208,
+					pointerSize: 134,
 				},
+				oid:      "9eb98c817f04b051b3bcca591bcd4e03cec88018",
 				path:     "tf_model.h5",
 				security: {
 					avScan: {
@@ -141,12 +139,11 @@ describe("listFiles", () => {
 			},
 			{
 				lastCommit: {
-					author: {
-						date: "2018-11-14T23:35:08.000Z",
-					},
-					id:      "2f07d813ca87c8c709147704c87210359ccf2309",
-					subject: "Update vocab.txt",
+					date:  "2018-11-14T23:35:08.000Z",
+					id:    "2f07d813ca87c8c709147704c87210359ccf2309",
+					title: "Update vocab.txt",
 				},
+				oid:      "fb140275c155a9c7c5a3b3e0e77a9e839594a938",
 				path:     "vocab.txt",
 				security: {
 					avScan: {
diff --git a/packages/hub/src/lib/list-files.ts b/packages/hub/src/lib/list-files.ts
@@ -2,6 +2,7 @@ import { HUB_URL } from "../consts";
 import { createApiError } from "../error";
 import type { Credentials, RepoId } from "../types";
 import type { ApiIndexTreeEntryData } from "../types/api";
+import { parseLinkHeader } from "../utils";
 
 export type ListFileEntry = ApiIndexTreeEntryData;
 
@@ -29,7 +30,7 @@ export async function* listFiles(params: {
 	}${params.path ? "/" + params.path : ""}${params.recursive ? "?recursive=true" : ""}`;
 
 	while (url) {
-		const res = await fetch(url, {
+		const res: Response = await fetch(url, {
 			headers: {
 				accept: "application/json",
 				...(params.credentials ? { Authorization: `Bearer ${params.credentials.accessToken}` } : undefined),
@@ -40,16 +41,14 @@ export async function* listFiles(params: {
 			throw createApiError(res);
 		}
 
-		const json: { items: ApiIndexTreeEntryData[]; nextUrl?: string } = await res.json();
+		const items: ApiIndexTreeEntryData[] = await res.json();
 
-		for (const item of json.items) {
+		for (const item of items) {
 			yield item;
 		}
 
-		url = json.nextUrl
-			? json.nextUrl?.startsWith(params.hubUrl ?? HUB_URL)
-				? json.nextUrl
-				: `${params.hubUrl ?? HUB_URL}${json.nextUrl}`
-			: undefined;
+		const linkHeader = res.headers.get("Link");
+
+		url = linkHeader ? parseLinkHeader(linkHeader).next : undefined;
 	}
 }
diff --git a/packages/hub/src/types/api.d.ts b/packages/hub/src/types/api.d.ts
@@ -210,8 +210,10 @@ export interface ApiIndexTreeEntryData {
 	size: number;
 	path: string;
 	lfs?: {
-		oid:  string;
-		size: number;
+		oid:         string;
+		size:        number;
+		/** Size of the raw pointer file, 100~200 bytes */
+		pointerSize: number;
 	};
 	lastCommit: {
 		author: {
diff --git a/packages/hub/src/utils/index.ts b/packages/hub/src/utils/index.ts
@@ -6,3 +6,4 @@ export { promisesQueueStreaming } from "./promisesQueueStreaming";
 export { promisesQueue } from "./promisesQueue";
 export { sha256 } from "./sha256";
 export { randomUUID } from "./randomUUID";
+export { parseLinkHeader } from "./parseLinkHeader";
diff --git a/packages/hub/src/utils/parseLinkHeader.ts b/packages/hub/src/utils/parseLinkHeader.ts
@@ -0,0 +1,8 @@
+/**
+ * Parse Link HTTP header, eg `<https://huggingface.co/api/datasets/bigscience/P3/tree/main?recursive=1&cursor=...>; rel="next"`
+ */
+export function parseLinkHeader(header: string): Record<string, string> {
+	const regex = /<(https?:[/][/][^>]+)>;\s+rel="([^"]+)"/g;
+
+	return Object.fromEntries([...header.matchAll(regex)].map(([, url, rel]) => [rel, url]));
+}
