File Content: create_numpy_pickle.py """ This script is used to generate test data for joblib/test/test_numpy_pickle.py """ import sys import re # nosetests needs to be able to import this module even when numpy is # not installed try: import numpy as np except ImportError: np = None import joblib def get_joblib_version(joblib_version=joblib.__version__): """Normalize joblib version by removing suffix. >>> get_joblib_version('0.8.4') '0.8.4' >>> get_joblib_version('0.8.4b1') '0.8.4' >>> get_joblib_version('0.9.dev0') '0.9' """ matches = [re.match(r'(\d+).*', each) for each in joblib_version.split('.')] return '.'.join([m.group(1) for m in matches if m is not None]) def write_test_pickle(to_pickle, args): kwargs = {} compress = args.compress method = args.method joblib_version = get_joblib_version() py_version = '{0[0]}{0[1]}'.format(sys.version_info) numpy_version = ''.join(np.__version__.split('.')[:2]) # The game here is to generate the right filename according to the options. body = '_compressed' if (compress and method == 'zlib') else '' if compress: if method == 'zlib': kwargs['compress'] = True extension = '.gz' else: kwargs['compress'] = (method, 3) extension = '.pkl.{0}'.format(method) if args.cache_size: kwargs['cache_size'] = 0 body += '_cache_size' else: extension = '.pkl' pickle_filename = 'joblib_{0}{1}_pickle_py{2}_np{3}{4}'.format( joblib_version, body, py_version, numpy_version, extension) try: joblib.dump(to_pickle, pickle_filename, **kwargs) except Exception as e: # With old python version (=< 3.3.), we can arrive there when # dumping compressed pickle with LzmaFile. print("Error: cannot generate file '{0}' with arguments '{1}'. " "Error was: {2}".format(pickle_filename, kwargs, e)) else: print("File '{0}' generated successfuly.".format(pickle_filename)) if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description="Joblib pickle data " "generator.") parser.add_argument('--cache_size', action="store_true", help="Force creation of companion numpy " "files for pickled arrays.") parser.add_argument('--compress', action="store_true", help="Generate compress pickles.") parser.add_argument('--method', type=str, default='zlib', choices=['zlib', 'gzip', 'bz2', 'xz', 'lzma'], help="Set compression method.") # We need to be specific about dtypes in particular endianness # because the pickles can be generated on one architecture and # the tests run on another one. See # https://github.com/joblib/joblib/issues/279. to_pickle = [np.arange(5, dtype=np.dtype('<i8')), np.arange(5, dtype=np.dtype('<f8')), np.array([1, 'abc', {'a': 1, 'b': 2}], dtype='O'), # all possible bytes as a byte string # .tostring actually returns bytes and is a # compatibility alias for .tobytes which was # added in 1.9.0 np.arange(256, dtype=np.uint8).tostring(), np.matrix([0, 1, 2], dtype=np.dtype('<i8')), # unicode string with non-ascii chars u"C'est l'\xe9t\xe9 !"] write_test_pickle(to_pickle, parser.parse_args()) Submit