Add seek and tell functions from so3g.G3IndexedReader

core/include/core/G3Reader.h

@@ -15,6 +15,8 @@ class G3Reader : public G3Module {
                  float timeout = -1.);
 
 	void Process(G3FramePtr frame, std::deque<G3FramePtr> &out);
+	int Seek(int offset);
+	int Tell();
 
 private:
 	void StartFile(std::string path);

core/src/G3Reader.cxx

@@ -10,7 +10,7 @@ G3Reader::G3Reader(std::string filename, int n_frames_to_read,
     n_frames_read_(0), timeout_(timeout)
 {
 	boost::filesystem::path fpath(filename);
-	if (filename.find("://") == -1 &&
+	if (filename.find("://") == std::string::npos &&
 	   (!boost::filesystem::exists(fpath) ||
 	    !boost::filesystem::is_regular_file(fpath)))
 		log_fatal("Could not find file %s", filename.c_str());
@@ -27,7 +27,7 @@ G3Reader::G3Reader(std::vector<std::string> filename, int n_frames_to_read,
 
 	for (auto i = filename.begin(); i != filename.end(); i++){
 		boost::filesystem::path fpath(*i);
-		if (i->find("://") == -1 &&
+		if (i->find("://") == std::string::npos &&
 		   (!boost::filesystem::exists(fpath) ||
 		    !boost::filesystem::is_regular_file(fpath)))
 			log_fatal("Could not find file %s", i->c_str());
@@ -108,6 +108,14 @@ void G3Reader::Process(G3FramePtr frame, std::deque<G3FramePtr> &out)
 	n_frames_read_++;
 }
 
+int G3Reader::Seek(int offset) {
+	return boost::iostreams::seek(stream_, offset, std::ios_base::beg);
+}
+
+int G3Reader::Tell() {
+	return boost::iostreams::seek(stream_, 0, std::ios_base::cur);
+}
+
 PYBINDINGS("core") {
 	using namespace boost::python;
 
@@ -118,14 +126,18 @@ PYBINDINGS("core") {
 	      "or an iterable of files to be read in sequence. If "
 	      "n_frames_to_read is greater than zero, will stop after "
 	      "n_frames_to_read frames rather than at the end of the file[s]. "
-              "The timeout parameter can used to enable socket timeout for tcp "
-              "streams, resulting in EOF behavior on expiry; unfortunately this "
-              "cannot be used for polling, you have to close the connection.",
-              init<std::string, int, float>((arg("filename"),
-                arg("n_frames_to_read")=0,arg("timeout")=-1.)))
-                .def(init<std::vector<std::string>, int, float>((arg("filename"), 
-                  arg("n_frames_to_read")=0, arg("timeout")=-1.)))
-		.def_readonly("__g3module__", true)
+	      "The timeout parameter can used to enable socket timeout for tcp "
+	      "streams, resulting in EOF behavior on expiry; unfortunately this "
+	      "cannot be used for polling, you have to close the connection. "
+	      "Use the `tell` and `seek` methods to record the position of and "
+	      "seek to the beginning of a particular frame in the file.",
+	init<std::string, int, float>((arg("filename"),
+	    arg("n_frames_to_read")=0,arg("timeout")=-1.)))
+	.def(init<std::vector<std::string>, int, float>((arg("filename"),
+	    arg("n_frames_to_read")=0, arg("timeout")=-1.)))
+	.def("tell", &G3Reader::Tell)
+	.def("seek", &G3Reader::Seek)
+	.def_readonly("__g3module__", true)
 	;
 }
 

core/tests/fileio.py

@@ -1,23 +1,34 @@
 #!/usr/bin/env python
 
 from spt3g import core
-import time, sys
+import time
 
 # File to disk
 pipe = core.G3Pipeline()
 pipe.Add(core.G3InfiniteSource, type=core.G3FrameType.Timepoint, n=10)
+# Drop a wiring frame in the middle
+count = 0
+def addwiring(fr):
+    global count
+    if fr.type == core.G3FrameType.Timepoint:
+        count += 1
+        if count == 6:
+            fr2 = core.G3Frame(core.G3FrameType.Wiring)
+            return [fr2, fr]
+pipe.Add(addwiring)
 n = 0
 def addinfo(fr):
-	global n
-	if fr.type != core.G3FrameType.Timepoint:
-		return
-	fr['time'] = core.G3Time(int(time.time()*core.G3Units.s))
-	fr['count'] = n
-	n += 1
+    global n
+    if fr.type != core.G3FrameType.Timepoint:
+        return
+    fr['time'] = core.G3Time(int(time.time()*core.G3Units.s))
+    fr['count'] = n
+    n += 1
 pipe.Add(addinfo)
 pipe.Add(core.Dump)
 pipe.Add(core.G3Writer, filename='test.g3')
 pipe.Run()
+assert n == 10, 'Wrong number of frames written (%d should be %d)' % (n, 10)
 
 # And back from disk
 print('Reading')
@@ -26,22 +37,67 @@ def addinfo(fr):
 pipe.Add(core.Dump)
 n = 0
 def checkinfo(fr):
-	global n
-	if fr.type != core.G3FrameType.Timepoint:
-		return
-	if 'time' not in fr:
-		print('No time key in frame')
-		sys.exit(1)
-	if fr['count'] != n:
-		print('Out of order frame')
-		sys.exit(1)
-	n += 1
+    global n
+    if fr.type != core.G3FrameType.Timepoint:
+        return
+    assert 'time' in fr, 'No time key in frame'
+    assert fr['count'] == n, 'Out of order frame'
+    n += 1
 pipe.Add(checkinfo)
 pipe.Run()
 
-if n != 10:
-	print('Wrong number of frames (%d should be %d)' % (n, 10))
-	sys.exit(1)
+assert n == 10, 'Wrong number of frames read (%d should be %d)' % (n, 10)
 
-sys.exit(0)
+# Indexing
+class CachingReader:
+    def __init__(self, filename='test.g3'):
+        self.reader = core.G3Reader(filename='test.g3')
+        self.w_pos = None
 
+    def __call__(self, frame):
+        assert frame is None
+        pos = self.reader.tell()
+        fr = self.reader(frame)
+        if not len(fr):
+            return fr
+        if fr[0].type == core.G3FrameType.Wiring:
+            self.w_pos = pos
+        return fr
+
+cacher = CachingReader()
+
+pipe = core.G3Pipeline()
+pipe.Add(cacher)
+pipe.Add(core.Dump)
+pipe.Run()
+
+assert cacher.w_pos is not None, 'Missing wiring frame'
+
+# Using cached index
+class CachedReader:
+    def __init__(self, filename='test.g3', start=None):
+        self.reader = core.G3Reader(filename=filename)
+        self.pos = start
+
+    def __call__(self, frame):
+        assert frame is None
+        if self.pos is not None:
+            self.reader.seek(self.pos)
+        if self.pos is not None:
+            assert self.reader.tell() == self.pos
+        fr = self.reader(frame)
+        if not len(fr):
+            return fr
+        if self.pos is not None:
+            assert fr[0].type == core.G3FrameType.Wiring
+            self.pos = None
+        return fr
+
+cached = CachedReader(start=cacher.w_pos)
+
+pipe = core.G3Pipeline()
+pipe.Add(cached)
+pipe.Add(core.Dump)
+pipe.Run()
+
+assert cached.pos is None, "Missing wiring frame"