Source code for pex.util

# Copyright 2014 Pants project contributors (see CONTRIBUTORS.md).# Licensed under the Apache License, Version 2.0 (see LICENSE).from__future__importabsolute_importimportcontextlibimporterrnoimportosimportshutilimportuuidfromhashlibimportsha1fromthreadingimportLockfrompkg_resourcesimportfind_distributions,resource_isdir,resource_listdir,resource_stringfrom.commonimportsafe_mkdir,safe_mkdtemp,safe_open,safe_rmtreefrom.findersimportregister_findersclassDistributionHelper(object):@classmethoddefwalk_data(cls,dist,path='/'):"""Yields filename, stream for files identified as data in the distribution"""forrel_fninfilter(None,dist.resource_listdir(path)):full_fn=os.path.join(path,rel_fn)ifdist.resource_isdir(full_fn):forfn,streamincls.walk_data(dist,full_fn):yieldfn,streamelse:yieldfull_fn[1:],dist.get_resource_stream(dist._provider,full_fn)@staticmethoddefzipsafe(dist):"""Returns whether or not we determine a distribution is zip-safe."""# zip-safety is only an attribute of eggs. wheels are considered never# zip safe per implications of PEP 427.ifhasattr(dist,'egg_info')anddist.egg_info.endswith('EGG-INFO'):egg_metadata=dist.metadata_listdir('')return'zip-safe'inegg_metadataand'native_libs.txt'notinegg_metadataelse:returnFalse@classmethoddefaccess_zipped_assets(cls,static_module_name,static_path,dir_location=None):""" Create a copy of static resource files as we can't serve them from within the pex file. :param static_module_name: Module name containing module to cache in a tempdir :type static_module_name: string, for example 'twitter.common.zookeeper' or similar :param static_path: Module name, for example 'serverset' :param dir_location: create a new temporary directory inside, or None to have one created :returns temp_dir: Temporary directory with the zipped assets inside :rtype: str """# asset_path is initially a module name that's the same as the static_path, but will be# changed to walk the directory treedefwalk_zipped_assets(static_module_name,static_path,asset_path,temp_dir):forassetinresource_listdir(static_module_name,asset_path):asset_target=os.path.normpath(os.path.join(os.path.relpath(asset_path,static_path),asset))ifresource_isdir(static_module_name,os.path.join(asset_path,asset)):safe_mkdir(os.path.join(temp_dir,asset_target))walk_zipped_assets(static_module_name,static_path,os.path.join(asset_path,asset),temp_dir)else:withopen(os.path.join(temp_dir,asset_target),'wb')asfp:path=os.path.join(static_path,asset_target)file_data=resource_string(static_module_name,path)fp.write(file_data)ifdir_locationisNone:temp_dir=safe_mkdtemp()else:temp_dir=dir_locationwalk_zipped_assets(static_module_name,static_path,static_path,temp_dir)returntemp_dir@classmethoddefdistribution_from_path(cls,path,name=None):"""Return a distribution from a path. If name is provided, find the distribution. If none is found matching the name, return None. If name is not provided and there is unambiguously a single distribution, return that distribution otherwise None. """# Monkeypatch pkg_resources finders should it not already be so.register_finders()ifnameisNone:distributions=list(find_distributions(path))iflen(distributions)==1:returndistributions[0]else:fordistinfind_distributions(path):ifdist.project_name==name:returndistclassCacheHelper(object):@classmethoddefupdate_hash(cls,filelike,digest):"""Update the digest of a single file in a memory-efficient manner."""block_size=digest.block_size*1024forchunkiniter(lambda:filelike.read(block_size),b''):digest.update(chunk)@classmethoddefhash(cls,path,digest=None,hasher=sha1):"""Return the digest of a single file in a memory-efficient manner."""ifdigestisNone:digest=hasher()withopen(path,'rb')asfh:cls.update_hash(fh,digest)returndigest.hexdigest()@classmethoddef_compute_hash(cls,names,stream_factory):digest=sha1()digest.update(''.join(names).encode('utf-8'))fornameinnames:withcontextlib.closing(stream_factory(name))asfp:cls.update_hash(fp,digest)returndigest.hexdigest()@classmethoddefzip_hash(cls,zf,prefix=''):"""Return the hash of the contents of a zipfile, comparable with a cls.dir_hash."""prefix_length=len(prefix)names=sorted(name[prefix_length:]fornameinzf.namelist()ifname.startswith(prefix)andnotname.endswith('.pyc')andnotname.endswith('/'))defstream_factory(name):returnzf.open(prefix+name)returncls._compute_hash(names,stream_factory)@classmethoddef_iter_files(cls,directory):normpath=os.path.normpath(directory)forroot,_,filesinos.walk(normpath):forfinfiles:yieldos.path.relpath(os.path.join(root,f),normpath)@classmethoddefpex_hash(cls,d):"""Return a reproducible hash of the contents of a directory."""names=sorted(fforfincls._iter_files(d)ifnot(f.endswith('.pyc')orf.startswith('.')))defstream_factory(name):returnopen(os.path.join(d,name),'rb')# noqa: T802returncls._compute_hash(names,stream_factory)@classmethoddefdir_hash(cls,d):"""Return a reproducible hash of the contents of a directory."""names=sorted(fforfincls._iter_files(d)ifnotf.endswith('.pyc'))defstream_factory(name):returnopen(os.path.join(d,name),'rb')# noqa: T802returncls._compute_hash(names,stream_factory)@classmethoddefcache_distribution(cls,zf,source,target_dir):"""Possibly cache an egg from within a zipfile into target_cache. Given a zipfile handle and a filename corresponding to an egg distribution within that zip, maybe write to the target cache and return a Distribution."""dependency_basename=os.path.basename(source)ifnotos.path.exists(target_dir):target_dir_tmp=target_dir+'.'+uuid.uuid4().hexfornameinzf.namelist():ifname.startswith(source)andnotname.endswith('/'):# strip off prefix + '/'target_name=os.path.join(dependency_basename,name[len(source)+1:])withcontextlib.closing(zf.open(name))aszi:withsafe_open(os.path.join(target_dir_tmp,target_name),'wb')asfp:shutil.copyfileobj(zi,fp)try:os.rename(target_dir_tmp,target_dir)exceptOSErrorase:ife.errno==errno.ENOTEMPTY:safe_rmtree(target_dir_tmp)else:raisedist=DistributionHelper.distribution_from_path(target_dir)assertdistisnotNone,'Failed to cache distribution %s'%sourcereturndist

[docs]classMemoizer(object):"""A thread safe class for memoizing the results of a computation."""def__init__(self):self._data={}self._lock=Lock()defget(self,key,default=None):withself._lock:returnself._data.get(key,default)defstore(self,key,value):withself._lock:self._data[key]=value