Возникла проблема с настраиваемым итератором, заключающаяся в том, что он выполняет итерацию по файлу только один раз. Я вызываю seek (0)
для соответствующего файлового объекта между итерациями, но StopIteration
вызывается при первом вызове next ()
при втором запуске через. Я чувствую, что упускаю из виду что-то очевидное, но я хотел бы увидеть это свежим взглядом:
class MappedIterator(object):
"""
Given an iterator of dicts or objects and a attribute mapping dict,
will make the objects accessible via the desired interface.
Currently it will only produce dictionaries with string values. Can be
made to support actual objects later on. Somehow... :D
"""
def __init__(self, obj=None, mapping={}, *args, **kwargs):
self._obj = obj
self._mapping = mapping
self.cnt = 0
def __iter__(self):
return self
def reset(self):
self.cnt = 0
def next(self):
try:
try:
item = self._obj.next()
except AttributeError:
item = self._obj[self.cnt]
# If no mapping is provided, an empty object will be returned.
mapped_obj = {}
for mapped_attr in self._mapping:
attr = mapped_attr.attribute
new_attr = mapped_attr.mapped_name
val = item.get(attr, '')
val = str(val).strip() # get rid of whitespace
# TODO: apply transformers...
# This allows multi attribute mapping or grouping of multiple
# attributes in to one.
try:
mapped_obj[new_attr] += val
except KeyError:
mapped_obj[new_attr] = val
self.cnt += 1
return mapped_obj
except (IndexError, StopIteration):
self.reset()
raise StopIteration
class CSVMapper(MappedIterator):
def __init__(self, reader, mapping={}, *args, **kwargs):
self._reader = reader
self._mapping = mapping
self._file = kwargs.pop('file')
super(CSVMapper, self).__init__(self._reader, self._mapping, *args, **kwargs)
@classmethod
def from_csv(cls, file, mapping, *args, **kwargs):
# TODO: Parse kwargs for various DictReader kwargs.
return cls(reader=DictReader(file), mapping=mapping, file=file)
def __len__(self):
return int(self._reader.line_num)
def reset(self):
if self._file:
self._file.seek(0)
super(CSVMapper, self).reset()
Пример использования:
file = open('somefile.csv', 'rb') # say this file has 2 rows + a header row
mapping = MyMappingClass() # this isn't really relevant
reader = CSVMapper.from_csv(file, mapping)
# > 'John'
# > 'Bob'
for r in reader:
print r['name']
# This won't print anything
for r in reader:
print r['name']