diff --git a/servicex_analysis_utils/dataset_resolver.py b/servicex_analysis_utils/dataset_resolver.py index 265cfb8..9157476 100644 --- a/servicex_analysis_utils/dataset_resolver.py +++ b/servicex_analysis_utils/dataset_resolver.py @@ -33,13 +33,14 @@ def ds_type_resolver( - ds_name: Union[str, list[str]], + ds_name: Union[str, list[str], int], ) -> Union[dataset.FileList, dataset.Rucio, dataset.XRootD, dataset.CERNOpenData]: """Determine the type of dataset based on the input string and then return the ServiceX dataset object. Args: - ds_name (str): Name of the dataset to fetch. + ds_name (str, int, list): Name of the dataset to fetch. Can be a URL, Rucio DID, + CERN Open Data ID (as str or int), EOS path, or list of files. Returns: dataset: The dataset object @@ -71,6 +72,9 @@ def find_scope(name): if isinstance(ds_name, list): return dataset.FileList(ds_name) + elif isinstance(ds_name, int): + return dataset.CERNOpenData(ds_name) + elif re.match(r"^https?://", ds_name): url = ds_name diff --git a/tests/test_dataset_resolver.py b/tests/test_dataset_resolver.py index 5a9edfd..066ab66 100644 --- a/tests/test_dataset_resolver.py +++ b/tests/test_dataset_resolver.py @@ -37,6 +37,7 @@ ("test:data", dataset.Rucio), ("rucio://test:test", dataset.Rucio), ("123", dataset.CERNOpenData), + (123, dataset.CERNOpenData), ("root://eosatlas.cern.ch//eos/", dataset.FileList), ("root://eosatlas.cern.ch//eos/*", dataset.XRootD), (["root://eosatlas.cern.ch//eos/", "https://test.com"], dataset.FileList),