18530227417

Committed 15 Oct 2025 01:17PM UTC coverage: 97.071% (-0.4%) from 97.425%

Build # 18530227417

Build Type

Pull #10

github

Committed by

web-flow

Commit Message

Merge 5c6b8dcdb into 3e8a9dfa3

Pull Request Pull Request #10: Let find() take format as argument to allow returning as list

Coverage Stats

10 of 11 new or added lines in 2 files covered. (90.91%)

232 of 239 relevant lines covered (97.07%)

0.97 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.39

/hepdata_cli/api.py

# -*- coding: utf-8 -*-

from .version import __version__
from .resilient_requests import resilient_requests

import tarfile
import sys
import re
import os
import errno

SITE_URL = "https://www.hepdata.net"
# SITE_URL = "http://127.0.0.1:5000"

UPLOAD_MAX_SIZE = 52000000  # Upload limit in bytes
ALLOWED_FORMATS = ['csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json']

MAX_MATCHES, MATCHES_PER_PAGE = (10000, 10) if "pytest" not in sys.modules else (144, 12)


class Client(object):
    """API class to handle all queries to HEPData."""

    def __init__(self, verbose=False):
        """
        Initialises the client object.

        :param verbose: prints additional output.
        """
        self.verbose = verbose
        self.version = __version__
        # check service availability
        resilient_requests('get', SITE_URL + '/ping')

    def find(self, query, keyword=None, ids=None, max_matches=MAX_MATCHES, matches_per_page=MATCHES_PER_PAGE, format=str):
        """
        Search function for the hepdata database. Calls hepdata.net search function.

        :param query: string passed to hepdata.net search function. See advanced search tips at hepdata.net.
        :param keyword: filters return dictionary for given keyword. Exact match is first attempted, otherwise partial match is accepted.
        :param ids: accepts one of ("arxiv", "inspire", "hepdata").
        :param max_matches: maximum number of matches to return. Default is 10,000.
        :param matches_per_page: number of matches per page. Default is 10.
        :param format: specifies the return format if 'ids' is specified. Allowed formats are: str, list, set, tuple. Default is str.

        :return: returns a list of (filtered if 'keyword' is specified) dictionaries for the search matches. If 'ids' is specified it instead returns a list of ids in the format 'format'.
        """
        find_results = []
        for counter in range(int(max_matches / matches_per_page)):
            counter += 1
            response = self._query(query, page=counter, size=matches_per_page)
            data = response.json()
            if len(data['results']) == 0:
                break
            elif keyword is None and ids is None:
                # return full list of dictionary
                find_results += data['results']
            else:
                assert ids in [None, "arxiv", "inspire", "hepdata", "id"], "allowed ids are: arxiv, inspire and hepdata"
                if ids is not None:
                    if ids == "hepdata":
                        ids = "id"
                    keyword = ids
                # return specific dictionary entry (exact match)
                if any([keyword in result.keys() for result in data['results']]):
                    if ids is None:
                        find_results += [{keyword: result[keyword]} for result in data['results'] if keyword in result.keys()]
                    else:
                        find_results += [str(result[keyword]).replace("arXiv:", "") for result in data['results'] if keyword in result.keys()]
                # return specific dictionary entry (partial match)
                elif any([any([keyword in key for key in result.keys()]) for result in data['results']]):
                    if ids is None:
                        find_results += [{key: result[key] for key in result.keys() if keyword in key} for result in data['results']]
                    else:
                        find_results += [[str(result[key]).replace("arXiv:", "") for key in result.keys() if keyword in key][0]
                                         if len([result[key] for key in result.keys() if keyword in key]) > 0 else "" for result in data['results']]
            if len(data['results']) < matches_per_page:
                break
        if ids is None:
            return find_results
        else:
            if format==str:
                return ' '.join(find_results)
            elif format==list:
                return find_results
            elif format in (set, tuple):
                return format(find_results)
            else:
                raise TypeError(f"Cannot return results in specfied format: {format}. Allowed formats are: {str}, {list}.")

    def download(self, id_list, file_format=None, ids=None, table_name='', download_dir='./hepdata-downloads'):
        """
        Downloads from the hepdata database the specified records.

        :param id_list: list of ids to download. These can be obtained by the find function.
        :param file_format: accepts one of ('csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json'). Specifies the download file format.
        :param ids: accepts one of ('inspire', 'hepdata'). It specifies what type of ids have been passed.
        :param table_name: restricts download to specific tables.
        :param download_dir: defaults to ./hepdata-downloads. Specifies where to download the files.

        :return: dictionary mapping id to list of downloaded files.
        :rtype: dict[int, list[str]]
        """

        url_map = self._build_urls(id_list, file_format, ids, table_name)
        file_map = {}
        for record_id, url in url_map.items():
            if self.verbose is True:
                print("Downloading: " + url)
            files_downloaded = download_url(url, download_dir)
            file_map[record_id] = files_downloaded
        return file_map

    def fetch_names(self, id_list, ids=None):
        """
        Returns the names of the tables in the provided records. These are the possible inputs of table_name parameter in download function.

        :param id_list: list of id of records of which to return table names.
        :param ids: accepts one of ('inspire', 'hepdata'). It specifies what type of ids have been passed.
        """
        url_map = self._build_urls(id_list, 'json', ids, '')
        table_names = []
        for url in url_map.values():
            response = resilient_requests('get', url)
            json_dict = response.json()
            table_names += [[data_table['name'] for data_table in json_dict['data_tables']]]
        return table_names

    def upload(self, path_to_file, email, recid=None, invitation_cookie=None, sandbox=True, password=None):
        """
        Upload record.

        :param path_to_file: path of file to be uploaded.
        :param email: email address of existing HEPData user.
        :recid: HEPData ID (not the INSPIRE ID) of an existing record.
        :invitation_cookie: token sent in the invitation email for a non-sandbox record.
        :sandbox: True (default) or False if the file should be uploaded to the sandbox.
        :password: password of existing HEPData user (prompt if not specified).
        """
        file_size = os.path.getsize(path_to_file)
        assert file_size < UPLOAD_MAX_SIZE,\
            '{} too large ({} bytes > {} bytes)'.format(path_to_file, file_size, UPLOAD_MAX_SIZE)
        files = {'hep_archive': open(path_to_file, 'rb')}
        data = {'email': email, 'recid': recid, 'invitation_cookie': invitation_cookie, 'sandbox': sandbox, 'pswd': password}
        resilient_requests('post', SITE_URL + '/record/cli_upload', data=data, files=files)
        # print upload location
        if sandbox is True and recid is None:
            print('Uploaded ' + path_to_file + ' to a new record at ' + SITE_URL + '/record/sandbox')
        elif sandbox is True and recid is not None:
            print('Uploaded ' + path_to_file + ' to ' + SITE_URL + '/record/sandbox/' + str(recid))
        else:
            print('Uploaded ' + path_to_file + ' to ' + SITE_URL + '/record/' + str(recid))

    def _build_urls(self, id_list, file_format, ids, table_name):
        """
        Builds urls for download and fetch_names, given the specified parameters.
        
        :param id_list: list of ids to download. Format is tuple, list, set or space-separated string.
        :param file_format: accepts one of ('csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json').
        :param ids: accepts one of ('inspire', 'hepdata').
        :param table_name: restricts download to specific tables.
        
        :return: dictionary mapping id to url.
        """
        if isinstance(id_list, str):
            id_list = id_list.split()
        assert len(id_list) > 0, 'Ids are required.'
        assert file_format in ALLOWED_FORMATS, f"allowed formats are: {ALLOWED_FORMATS}"
        assert ids in ['inspire', 'hepdata'], "allowed ids are: inspire and hepdata."
        if table_name == '':
            params = {'format': file_format}
        else:
            params = {'format': file_format, 'table': table_name}
        url_mapping = {}
        for id_entry in id_list:
            url = resilient_requests('get', SITE_URL + '/record/' + ('ins' if ids == 'inspire' else '') + id_entry, params=params).url.replace('%2525', '%25')
            url_mapping[id_entry] = url
        # TODO: Investigate root cause of double URL encoding (https://github.com/HEPData/hepdata-cli/issues/8).
        return url_mapping

    def _query(self, query, page, size):
        """Builds the search query passed to hepdata.net."""
        url = SITE_URL + '/search/?q=' + query + '&format=json&page=' + str(page) + '&size=' + str(size)
        response = resilient_requests('get', url)
        if self.verbose is True:
            print('Looking up: ' + url)
        return response


def mkdir(directory):
    if not os.path.exists(directory):
        try:
            os.makedirs(directory)
        except OSError as exc:   # Guard against race condition (directory created between os.path.exists and os.makedirs)
            if exc.errno != errno.EEXIST:
                raise Exception


def download_url(url, download_dir):
    """Download file and if necessary extract it."""
    files_downloaded = []
    assert is_downloadable(url), "Given url is not downloadable: {}".format(url)
    response = resilient_requests('get', url, allow_redirects=True)
    if url[-4:] == 'json':
        filename = 'HEPData-' + url.split('/')[-1].split("?")[0] + ".json"
    else:
        filename = getFilename_fromCd(response.headers.get('content-disposition'))
    if filename[0] == '"' and filename[-1] == '"':
        filename = filename[1:-1]
    filepath = download_dir + "/" + filename
    mkdir(os.path.dirname(filepath))
    open(filepath, 'wb').write(response.content)
    if filepath.endswith("tar.gz") or filepath.endswith("tar"):
        tar = None
        try:
            tar = tarfile.open(filepath, "r:gz" if filepath.endswith("tar.gz") else "r:")
            extract_dir = os.path.abspath(os.path.dirname(filepath))
            tar.extractall(path=os.path.dirname(filepath))
            for member in tar.getmembers():
                if member.isfile():
                    extracted_path = os.path.join(os.path.dirname(filepath), member.name)
                    abs_extracted_path = os.path.abspath(extracted_path)
                    if abs_extracted_path.startswith(extract_dir + os.sep) and os.path.exists(abs_extracted_path):
                        files_downloaded.append(abs_extracted_path)
                    elif not abs_extracted_path.startswith(extract_dir + os.sep):
                        raise ValueError(f"Attempted path traversal for file {member.name}")
                    else:
                        raise FileNotFoundError(f"Extracted file {member.name} not found")
        except Exception as e:
            raise Exception(f"Failed to extract {filepath}: {str(e)}")
        finally:
            if tar:
                tar.close()
            if os.path.exists(filepath):
                os.remove(filepath)
    else:
        files_downloaded.append(filepath)
    return files_downloaded


def getFilename_fromCd(cd):
    """Get filename from content-disposition."""
    if not cd:
        return None
    fname = re.findall('filename=(.+)', cd)
    if len(fname) == 0:
        return None
    return fname[0]


def is_downloadable(url):
    """Does the url contain a downloadable resource?"""
    header = resilient_requests('head', url, allow_redirects=True).headers
    content_type = header.get('content-type')
    if 'html' in content_type.lower():
        return False
    return True

1	# -- coding: utf-8 --
2
3	from .version import __version__	1✔
4	from .resilient_requests import resilient_requests	1✔
5
6	import tarfile	1✔
7	import sys	1✔
8	import re	1✔
9	import os	1✔
10	import errno	1✔
11
12	SITE_URL = "https://www.hepdata.net"	1✔
13	# SITE_URL = "http://127.0.0.1:5000"
14
15	UPLOAD_MAX_SIZE = 52000000 # Upload limit in bytes	1✔
16	ALLOWED_FORMATS = ['csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json']	1✔
17
18	MAX_MATCHES, MATCHES_PER_PAGE = (10000, 10) if "pytest" not in sys.modules else (144, 12)	1✔
19
20
21	class Client(object):	1✔
22	"""API class to handle all queries to HEPData."""
23
24	def __init__(self, verbose=False):	1✔
25	"""
26	Initialises the client object.
27
28	:param verbose: prints additional output.
29	"""
30	self.verbose = verbose	1✔
31	self.version = __version__	1✔
32	# check service availability
33	resilient_requests('get', SITE_URL + '/ping')	1✔
34
35	def find(self, query, keyword=None, ids=None, max_matches=MAX_MATCHES, matches_per_page=MATCHES_PER_PAGE, format=str):	1✔
36	"""
37	Search function for the hepdata database. Calls hepdata.net search function.
38
39	:param query: string passed to hepdata.net search function. See advanced search tips at hepdata.net.
40	:param keyword: filters return dictionary for given keyword. Exact match is first attempted, otherwise partial match is accepted.
41	:param ids: accepts one of ("arxiv", "inspire", "hepdata").
42	:param max_matches: maximum number of matches to return. Default is 10,000.
43	:param matches_per_page: number of matches per page. Default is 10.
44	:param format: specifies the return format if 'ids' is specified. Allowed formats are: str, list, set, tuple. Default is str.
45
46	:return: returns a list of (filtered if 'keyword' is specified) dictionaries for the search matches. If 'ids' is specified it instead returns a list of ids in the format 'format'.
47	"""
48	find_results = []	1✔
49	for counter in range(int(max_matches / matches_per_page)):	1✔
50	counter += 1	1✔
51	response = self._query(query, page=counter, size=matches_per_page)	1✔
52	data = response.json()	1✔
53	if len(data['results']) == 0:	1✔
54	break	1✔
55	elif keyword is None and ids is None:	1✔
56	# return full list of dictionary
57	find_results += data['results']	1✔
58	else:
59	assert ids in [None, "arxiv", "inspire", "hepdata", "id"], "allowed ids are: arxiv, inspire and hepdata"	1✔
60	if ids is not None:	1✔
61	if ids == "hepdata":	1✔
62	ids = "id"	1✔
63	keyword = ids	1✔
64	# return specific dictionary entry (exact match)
65	if any([keyword in result.keys() for result in data['results']]):	1✔
66	if ids is None:	1✔
67	find_results += [{keyword: result[keyword]} for result in data['results'] if keyword in result.keys()]	1✔
68	else:
69	find_results += [str(result[keyword]).replace("arXiv:", "") for result in data['results'] if keyword in result.keys()]	1✔
70	# return specific dictionary entry (partial match)
71	elif any([any([keyword in key for key in result.keys()]) for result in data['results']]):	1✔
72	if ids is None:	1✔
73	find_results += [{key: result[key] for key in result.keys() if keyword in key} for result in data['results']]	1✔
74	else:
75	find_results += [[str(result[key]).replace("arXiv:", "") for key in result.keys() if keyword in key][0]	1✔
76	if len([result[key] for key in result.keys() if keyword in key]) > 0 else "" for result in data['results']]
77	if len(data['results']) < matches_per_page:	1✔
78	break	1✔
79	if ids is None:	1✔
80	return find_results	1✔
81	else:
82	if format==str:	1✔
83	return ' '.join(find_results)	1✔
84	elif format==list:	1✔
85	return find_results	1✔
86	elif format in (set, tuple):	1✔
87	return format(find_results)	1✔
88	else:
NEW 89	raise TypeError(f"Cannot return results in specfied format: {format}. Allowed formats are: {str}, {list}.")	×
90
91	def download(self, id_list, file_format=None, ids=None, table_name='', download_dir='./hepdata-downloads'):	1✔
92	"""
93	Downloads from the hepdata database the specified records.
94
95	:param id_list: list of ids to download. These can be obtained by the find function.
96	:param file_format: accepts one of ('csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json'). Specifies the download file format.
97	:param ids: accepts one of ('inspire', 'hepdata'). It specifies what type of ids have been passed.
98	:param table_name: restricts download to specific tables.
99	:param download_dir: defaults to ./hepdata-downloads. Specifies where to download the files.
100
101	:return: dictionary mapping id to list of downloaded files.
102	:rtype: dict[int, list[str]]
103	"""
104
105	url_map = self._build_urls(id_list, file_format, ids, table_name)	1✔
106	file_map = {}	1✔
107	for record_id, url in url_map.items():	1✔
108	if self.verbose is True:	1✔
109	print("Downloading: " + url)	1✔
110	files_downloaded = download_url(url, download_dir)	1✔
111	file_map[record_id] = files_downloaded	1✔
112	return file_map	1✔
113
114	def fetch_names(self, id_list, ids=None):	1✔
115	"""
116	Returns the names of the tables in the provided records. These are the possible inputs of table_name parameter in download function.
117
118	:param id_list: list of id of records of which to return table names.
119	:param ids: accepts one of ('inspire', 'hepdata'). It specifies what type of ids have been passed.
120	"""
121	url_map = self._build_urls(id_list, 'json', ids, '')	1✔
122	table_names = []	1✔
123	for url in url_map.values():	1✔
124	response = resilient_requests('get', url)	1✔
125	json_dict = response.json()	1✔
126	table_names += [[data_table['name'] for data_table in json_dict['data_tables']]]	1✔
127	return table_names	1✔
128
129	def upload(self, path_to_file, email, recid=None, invitation_cookie=None, sandbox=True, password=None):	1✔
130	"""
131	Upload record.
132
133	:param path_to_file: path of file to be uploaded.
134	:param email: email address of existing HEPData user.
135	:recid: HEPData ID (not the INSPIRE ID) of an existing record.
136	:invitation_cookie: token sent in the invitation email for a non-sandbox record.
137	:sandbox: True (default) or False if the file should be uploaded to the sandbox.
138	:password: password of existing HEPData user (prompt if not specified).
139	"""
140	file_size = os.path.getsize(path_to_file)	1✔
141	assert file_size < UPLOAD_MAX_SIZE,\	1✔
142	'{} too large ({} bytes > {} bytes)'.format(path_to_file, file_size, UPLOAD_MAX_SIZE)
143	files = {'hep_archive': open(path_to_file, 'rb')}	1✔
144	data = {'email': email, 'recid': recid, 'invitation_cookie': invitation_cookie, 'sandbox': sandbox, 'pswd': password}	1✔
145	resilient_requests('post', SITE_URL + '/record/cli_upload', data=data, files=files)	1✔
146	# print upload location
147	if sandbox is True and recid is None:	1✔
148	print('Uploaded ' + path_to_file + ' to a new record at ' + SITE_URL + '/record/sandbox')	1✔
149	elif sandbox is True and recid is not None:	1✔
150	print('Uploaded ' + path_to_file + ' to ' + SITE_URL + '/record/sandbox/' + str(recid))	1✔
151	else:
152	print('Uploaded ' + path_to_file + ' to ' + SITE_URL + '/record/' + str(recid))	1✔
153
154	def _build_urls(self, id_list, file_format, ids, table_name):	1✔
155	"""
156	Builds urls for download and fetch_names, given the specified parameters.
157
158	:param id_list: list of ids to download. Format is tuple, list, set or space-separated string.
159	:param file_format: accepts one of ('csv', 'root', 'yaml', 'yoda', 'yoda1', 'yoda.h5', 'json').
160	:param ids: accepts one of ('inspire', 'hepdata').
161	:param table_name: restricts download to specific tables.
162
163	:return: dictionary mapping id to url.
164	"""
165	if isinstance(id_list, str):	1✔
166	id_list = id_list.split()	1✔
167	assert len(id_list) > 0, 'Ids are required.'	1✔
168	assert file_format in ALLOWED_FORMATS, f"allowed formats are: {ALLOWED_FORMATS}"	1✔
169	assert ids in ['inspire', 'hepdata'], "allowed ids are: inspire and hepdata."	1✔
170	if table_name == '':	1✔
171	params = {'format': file_format}	1✔
172	else:
173	params = {'format': file_format, 'table': table_name}	1✔
174	url_mapping = {}	1✔
175	for id_entry in id_list:	1✔
176	url = resilient_requests('get', SITE_URL + '/record/' + ('ins' if ids == 'inspire' else '') + id_entry, params=params).url.replace('%2525', '%25')	1✔
177	url_mapping[id_entry] = url	1✔
178	# TODO: Investigate root cause of double URL encoding (https://github.com/HEPData/hepdata-cli/issues/8).
179	return url_mapping	1✔
180
181	def _query(self, query, page, size):	1✔
182	"""Builds the search query passed to hepdata.net."""
183	url = SITE_URL + '/search/?q=' + query + '&format=json&page=' + str(page) + '&size=' + str(size)	1✔
184	response = resilient_requests('get', url)	1✔
185	if self.verbose is True:	1✔
186	print('Looking up: ' + url)	1✔
187	return response	1✔
188
189
190	def mkdir(directory):	1✔
191	if not os.path.exists(directory):	1✔
192	try:	1✔
193	os.makedirs(directory)	1✔
194	except OSError as exc: # Guard against race condition (directory created between os.path.exists and os.makedirs)	×
195	if exc.errno != errno.EEXIST:	×
196	raise Exception	×
197
198
199	def download_url(url, download_dir):	1✔
200	"""Download file and if necessary extract it."""
201	files_downloaded = []	1✔
202	assert is_downloadable(url), "Given url is not downloadable: {}".format(url)	1✔
203	response = resilient_requests('get', url, allow_redirects=True)	1✔
204	if url[-4:] == 'json':	1✔
205	filename = 'HEPData-' + url.split('/')[-1].split("?")[0] + ".json"	1✔
206	else:
207	filename = getFilename_fromCd(response.headers.get('content-disposition'))	1✔
208	if filename[0] == '"' and filename[-1] == '"':	1✔
209	filename = filename[1:-1]	1✔
210	filepath = download_dir + "/" + filename	1✔
211	mkdir(os.path.dirname(filepath))	1✔
212	open(filepath, 'wb').write(response.content)	1✔
213	if filepath.endswith("tar.gz") or filepath.endswith("tar"):	1✔
214	tar = None	1✔
215	try:	1✔
216	tar = tarfile.open(filepath, "r:gz" if filepath.endswith("tar.gz") else "r:")	1✔
217	extract_dir = os.path.abspath(os.path.dirname(filepath))	1✔
218	tar.extractall(path=os.path.dirname(filepath))	1✔
219	for member in tar.getmembers():	1✔
220	if member.isfile():	1✔
221	extracted_path = os.path.join(os.path.dirname(filepath), member.name)	1✔
222	abs_extracted_path = os.path.abspath(extracted_path)	1✔
223	if abs_extracted_path.startswith(extract_dir + os.sep) and os.path.exists(abs_extracted_path):	1✔
224	files_downloaded.append(abs_extracted_path)	1✔
225	elif not abs_extracted_path.startswith(extract_dir + os.sep):	1✔
226	raise ValueError(f"Attempted path traversal for file {member.name}")	1✔
227	else:
228	raise FileNotFoundError(f"Extracted file {member.name} not found")	1✔
229	except Exception as e:	1✔
230	raise Exception(f"Failed to extract {filepath}: {str(e)}")	1✔
231	finally:
232	if tar:	1✔
233	tar.close()	1✔
234	if os.path.exists(filepath):	1✔
235	os.remove(filepath)	1✔
236	else:
237	files_downloaded.append(filepath)	1✔
238	return files_downloaded	1✔
239
240
241	def getFilename_fromCd(cd):	1✔
242	"""Get filename from content-disposition."""
243	if not cd:	1✔
244	return None	1✔
245	fname = re.findall('filename=(.+)', cd)	1✔
246	if len(fname) == 0:	1✔
247	return None	1✔
248	return fname[0]	1✔
249
250
251	def is_downloadable(url):	1✔
252	"""Does the url contain a downloadable resource?"""
253	header = resilient_requests('head', url, allow_redirects=True).headers	1✔
254	content_type = header.get('content-type')	1✔
255	if 'html' in content_type.lower():	1✔
256	return False	1✔
257	return True	1✔

HEPData / hepdata-cli / 18530227417

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous