Iterator doesn't work well with keyboard interupt

[kain88-de]

Expected behaviour

When I break a loop like this with KeyboardInterupt I expect the trajectory to point to the same frame as at the beginning.

>>> start_frame = long_sim.trajectory.frame
>>> for ts in long_sim.trajectory:
>>>     s += ts.frame
# Interupt this loop
>>> long_sim.trajectory.frame == start_frame
True

This is very important when I have an analysis that extracts from information from mobile in the initialization. Otherwise code like this would return different results based on the following.

>>> class Analysis(AnalysisBase):
>>>     def __init__(self, ...):
>>>          super(...)__init__(...)
>>>           print(trajectory.frame)

>>> ana = Analysis(u.trajectory, mobile)
0
>>> ana.run()
# interrupt because I want a progressmeter since this runs for ages

>>> ana2 = Analysis(u.trajectory, mobile, quiet=False)
123 # the frame at which the interrupt was raised.
>>> ana2.run()

Actual behaviour

The last comparison returns false

Code to reproduce the behaviour

Gist

Currently version of MDAnalysis:

(run python -c "import MDAnalysis as mda; print(mda.__version__)") 0.16.0dev

[richardjgowers]

I think the frame after an iteration is actually pretty badly defined currently...

In [4]: from MDAnalysisTests.datafiles import PSF, DCD

In [5]: u = mda.Universe(PSF, DCD)

In [6]: print u.trajectory.ts.frame
0

In [7]: u.trajectory[10]
Out[7]: < Timestep 10 with unit cell dimensions [  0.   0.   0.  90.  90.  90.] >

In [8]: print u.trajectory.ts.frame
10

In [9]: for ts in u.trajectory[:]:
   ...:     pass
   ...: 

In [10]: print u.trajectory.ts.frame
0

I think we're blindly rewinding after an iteration rather than reverting to the previous frame.

I think for this issue, we'd have to make the iterator into a context manager to allow it to clean up even after a KeyboardInterrupt? Not sure if that's necessary though

[kain88-de]

Does the context manager work when calling an iterator?

[kain88-de]

I think the problem is that __iter__ returns the object itself. But that is hard to fix the way AtomGroups and TimeSteps depend on each other.

[richardjgowers]

I think the current context manager is off coordinates.base.IOBase and is more concerned with file handles, so you could do:

with mda.coordinates.XTCReader('this.xtc') as f:
    for ts in f:

[kain88-de]

second gist. I have the problem running a special analysis class. I don't do any interrupts here. Its strange that this doesn't work as expected. It looks like we don't to a good reset to frame 0 after the iteration has finished for an arbitrary iteration stop.

[richardjgowers]

Hmm, it might be because rewind is only called on StopIteration which we only hit at the end of a trajectory, not when we voluntarily stop.

[kain88-de]

I'll try a quick fix.

[kain88-de]

@jbarnoud has this been fixed with your recent updates to the iteration?

[jbarnoud]

@kain88-de I just tried and it does not solves it. I maintained the same behavior.

[richardjgowers]

I'm not entirely sure we're responsible for keyboard interrupts? If we wanted to be though, we'd have to put in an except KeyboardInterrupt inside the classes @jbarnoud made, and handle it there