Created
February 5, 2014 12:29
-
-
Save waaaaargh/8822627 to your computer and use it in GitHub Desktop.
Monkeypatch Stem's DirectoryReader to just care about the relays we're interested in.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import tarfile | |
import stem | |
from stem.descriptor.reader import DescriptorReader | |
class MyDescriptorReader(DescriptorReader): | |
def __init__(self, interesting_fingerprints, *args, **kwargs): | |
DescriptorReader.__init__(self, *args, **kwargs) | |
self._interesting_fingerprints = interesting_fingerprints | |
def _handle_archive(self, target): | |
# TODO: This would be nicer via the 'with' keyword, but tarfile's __exit__ | |
# method was added sometime after python 2.5. We should change this when | |
# we drop python 2.5 support. | |
tar_file = None | |
try: | |
self._notify_read_listeners(target) | |
tar_file = tarfile.open(target) | |
for tar_entry in [tar_entry for tar_entry in tar_file if tar_entry.name.split("/")[-1] in self._interesting_fingerprints]: | |
if tar_entry.isfile(): | |
entry = tar_file.extractfile(tar_entry) | |
try: | |
for desc in stem.descriptor.parse_file(entry, validate = self._validate, document_handler = self._document_handler, **self._kwargs): | |
if self._is_stopped.is_set(): | |
return | |
desc._set_path(os.path.abspath(target)) | |
desc._set_archive_path(entry.name) | |
self._unreturned_descriptors.put(desc) | |
self._iter_notice.set() | |
except TypeError as exc: | |
self._notify_skip_listeners(target, ParsingFailure(exc)) | |
except ValueError as exc: | |
self._notify_skip_listeners(target, ParsingFailure(exc)) | |
finally: | |
entry.close() | |
except IOError as exc: | |
self._notify_skip_listeners(target, ReadFailed(exc)) | |
finally: | |
if tar_file: | |
tar_file.close() | |
my_descriptors = [ | |
'extra-infos-2014-02.tar.bz2' | |
] | |
my_fingerprints = [ | |
'446691dfdbafbb4e5f021b5a2b73b5c8fc297b91' | |
] | |
# prints the contents of all the descriptor files | |
with MyDescriptorReader(my_fingerprints, my_descriptors) as reader: | |
for descriptor in reader: | |
print "Relay %s wrote %i and read %i bytes." % (descriptor.fingerprint, sum(descriptor.read_history_values), sum(descriptor.write_history_values)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment