Spaces:
Sleeping
Sleeping
File size: 3,669 Bytes
19e88d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | import { HUB_URL } from "../consts";
import { createApiError } from "../error";
import type { ApiDatasetInfo } from "../types/api/api-dataset";
import type { CredentialsParams } from "../types/public";
import { checkCredentials } from "../utils/checkCredentials";
import { parseLinkHeader } from "../utils/parseLinkHeader";
import { pick } from "../utils/pick";
export const DATASET_EXPAND_KEYS = [
"private",
"downloads",
"gated",
"likes",
"lastModified",
] as const satisfies readonly (keyof ApiDatasetInfo)[];
export const DATASET_EXPANDABLE_KEYS = [
"author",
"cardData",
"citation",
"createdAt",
"disabled",
"description",
"downloads",
"downloadsAllTime",
"gated",
"gitalyUid",
"lastModified",
"likes",
"paperswithcode_id",
"private",
// "siblings",
"sha",
"tags",
] as const satisfies readonly (keyof ApiDatasetInfo)[];
export interface DatasetEntry {
id: string;
name: string;
private: boolean;
downloads: number;
gated: false | "auto" | "manual";
likes: number;
updatedAt: Date;
}
export async function* listDatasets<
const T extends Exclude<(typeof DATASET_EXPANDABLE_KEYS)[number], (typeof DATASET_EXPAND_KEYS)[number]> = never,
>(
params?: {
search?: {
/**
* Will search in the dataset name for matches
*/
query?: string;
owner?: string;
tags?: string[];
};
hubUrl?: string;
additionalFields?: T[];
/**
* Set to limit the number of datasets returned.
*/
limit?: number;
/**
* Sort datasets by a specific field.
*/
sort?:
| "createdAt"
| "downloads"
| "likes"
| "lastModified"
| "likes30d"
| "trendingScore"
| "datasetsServerInfo.numRows"
| "mainSize"
| "id";
/**
* Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
*/
fetch?: typeof fetch;
} & Partial<CredentialsParams>,
): AsyncGenerator<DatasetEntry & Pick<ApiDatasetInfo, T>> {
const accessToken = params && checkCredentials(params);
let totalToFetch = params?.limit ?? Infinity;
const search = new URLSearchParams([
...Object.entries({
limit: String(Math.min(totalToFetch, 500)),
...(params?.search?.owner ? { author: params.search.owner } : undefined),
...(params?.search?.query ? { search: params.search.query } : undefined),
...(params?.sort ? { sort: params.sort } : undefined),
}),
...(params?.search?.tags?.map((tag) => ["filter", tag]) ?? []),
...DATASET_EXPAND_KEYS.map((val) => ["expand", val] satisfies [string, string]),
...(params?.additionalFields?.map((val) => ["expand", val] satisfies [string, string]) ?? []),
]).toString();
let url: string | undefined = `${params?.hubUrl || HUB_URL}/api/datasets` + (search ? "?" + search : "");
while (url) {
const res: Response = await (params?.fetch ?? fetch)(url, {
headers: {
accept: "application/json",
...(accessToken ? { Authorization: `Bearer ${accessToken}` } : undefined),
},
});
if (!res.ok) {
throw await createApiError(res);
}
const items: ApiDatasetInfo[] = await res.json();
for (const item of items) {
yield {
...(params?.additionalFields && pick(item, params.additionalFields)),
id: item._id,
name: item.id,
private: item.private,
downloads: item.downloads,
likes: item.likes,
gated: item.gated,
updatedAt: new Date(item.lastModified),
} as DatasetEntry & Pick<ApiDatasetInfo, T>;
totalToFetch--;
if (totalToFetch <= 0) {
return;
}
}
const linkHeader = res.headers.get("Link");
url = linkHeader ? parseLinkHeader(linkHeader).next : undefined;
// Could update limit in url to fetch less items if not all items of next page are needed.
}
}
|