💥 Do not fetch last commit by default in listFiles (#151)

coyotte508 · web-flow · commit ef71d2b7558f · 2023-04-20T00:57:03.000+02:00
diff --git a/packages/hub/src/lib/list-files.spec.ts b/packages/hub/src/lib/list-files.spec.ts
@@ -15,6 +15,67 @@ describe("listFiles", () => {
 
 		const files: ListFileEntry[] = [];
 
+		for await (const entry of cursor) {
+			files.push(entry);
+		}
+
+		assert.deepStrictEqual(files, [
+			{
+				oid: "dc08351d4dc0732d9c8af04070ced089b201ce2f",
+				path: ".gitattributes",
+				size: 345,
+				type: "file",
+			},
+			{
+				oid: "fca794a5f07ff8f963fe8b61e3694b0fb7f955df",
+				path: "config.json",
+				size: 313,
+				type: "file",
+			},
+			{
+				lfs: {
+					oid: "097417381d6c7230bd9e3557456d726de6e83245ec8b24f529f60198a67b203a",
+					size: 440473133,
+					pointerSize: 134,
+				},
+				oid: "ba5d19791be1dd7992e33bd61f20207b0f7f50a5",
+				path: "pytorch_model.bin",
+				size: 440473133,
+				type: "file",
+			},
+			{
+				lfs: {
+					oid: "a7a17d6d844b5de815ccab5f42cad6d24496db3850a2a43d8258221018ce87d2",
+					size: 536063208,
+					pointerSize: 134,
+				},
+				oid: "9eb98c817f04b051b3bcca591bcd4e03cec88018",
+				path: "tf_model.h5",
+				size: 536063208,
+				type: "file",
+			},
+			{
+				oid: "fb140275c155a9c7c5a3b3e0e77a9e839594a938",
+				path: "vocab.txt",
+				size: 231508,
+				type: "file",
+			},
+		]);
+	});
+
+	it("should fetch the list of files from the repo, including last commit", async () => {
+		const cursor = listFiles({
+			repo: {
+				name: "bert-base-uncased",
+				type: "model",
+			},
+			revision: "dd4bc8b21efa05ec961e3efc4ee5e3832a3679c7",
+			expand: true,
+			hubUrl: "https://huggingface.co",
+		});
+
+		const files: ListFileEntry[] = [];
+
 		for await (const entry of cursor) {
 			delete entry.security; // flaky
 			files.push(entry);
diff --git a/packages/hub/src/lib/list-files.ts b/packages/hub/src/lib/list-files.ts
@@ -17,11 +17,17 @@ export interface ListFileEntry {
 		/** Size of the raw pointer file, 100~200 bytes */
 		pointerSize: number;
 	};
-	lastCommit: {
+	/**
+	 * Only fetched if `expand` is set to `true` in the `listFiles` call.
+	 */
+	lastCommit?: {
 		date: string;
 		id: string;
 		title: string;
-	} | null;
+	};
+	/**
+	 * Only fetched if `expand` is set to `true` in the `listFiles` call.
+	 */
 	security?: unknown;
 }
 
@@ -40,6 +46,10 @@ export async function* listFiles(params: {
 	 * files in the repo.
 	 */
 	path?: string;
+	/**
+	 * Fetch `lastCommit` and `securityStatus` for each file.
+	 */
+	expand?: boolean;
 	revision?: string;
 	credentials?: Credentials;
 	hubUrl?: string;
@@ -48,7 +58,7 @@ export async function* listFiles(params: {
 	const repoId = toRepoId(params.repo);
 	let url: string | undefined = `${params.hubUrl || HUB_URL}/api/${repoId.type}s/${repoId.name}/tree/${
 		params.revision || "main"
-	}${params.path ? "/" + params.path : ""}${params.recursive ? "?recursive=true" : ""}`;
+	}${params.path ? "/" + params.path : ""}?recursive=${!!params.recursive}&expand=${!!params.expand}`;
 
 	while (url) {
 		const res: Response = await fetch(url, {
diff --git a/packages/hub/src/types/api/api-index-tree.d.ts b/packages/hub/src/types/api/api-index-tree.d.ts
@@ -9,11 +9,11 @@ export interface ApiIndexTreeEntry {
 		/** Size of the raw pointer file, 100~200 bytes */
 		pointerSize: number;
 	};
-	lastCommit: {
+	lastCommit?: {
 		date: string;
 		id: string;
 		title: string;
-	} | null;
+	};
 	security?: ApiFileScanResult;
 }
 
