add saveCubemap(); robustify screenshot methods

direct/src/showbase/ShowBase.py

@@ -1530,59 +1530,131 @@ class ShowBase(DirectObject.DirectObject):
         if self.camFrustumVis:
             self.camFrustumVis.removeNode()
 
-    def screenshot(self, namePrefix='screenshot'):
-        filename = self.win.saveScreenshotDefault(namePrefix)
-        if filename.empty():
-            # The screenshot attempt failed for some reason.
-            return 0
+    def screenshot(self, namePrefix = 'screenshot',
+                   defaultFilename = 1, source = None):
+        """ Captures a screenshot from the main window or from the
+        specified window or Texture and writes it to a filename in the
+        current directory (or to a specified directory).
+
+        If defaultFilename is True, the filename is synthesized by
+        appending namePrefix to a default filename suffix (including
+        the filename extension) specified in the Config variable
+        screenshot-filename.  Otherwise, if defaultFilename is False,
+        the entire namePrefix is taken to be the filename to write,
+        and this string should include a suitable filename extension
+        that will be used to determine the type of image file to
+        write.
+
+        Normally, the source is a GraphicsWindow, GraphicsBuffer or
+        DisplayRegion.  If a Texture is supplied instead, it must have
+        a ram image (that is, if it was generated by
+        makeTextureBuffer() or makeCubeMap(), the parameter toRam
+        should have been set true).  If it is a cube map texture as
+        generated by makeCubeMap(), namePrefix should contain the hash
+        mark ('#') character. """
+        
+        if source == None:
+            source = self.win
+
+        if defaultFilename:
+            filename = GraphicsOutput.makeScreenshotFilename(namePrefix)
+        else:
+            filename = Filename(namePrefix)
+
+        if isinstance(source, Texture):
+            if source.getZSize() > 1:
+                saved = source.writePages(filename)
+            else:
+                saved = source.write(filename)
+        else:
+            saved = source.saveScreenshot(filename)
 
         # Announce to anybody that a screenshot has been taken
         messenger.send('screenshot', [filename])
-        return 1
+
+        return saved
+
+    def saveCubemap(self, namePrefix = 'cubemap_#.png',
+                    defaultFilename = 0, source = None,
+                    camera = None, size = 128,
+                    cameraMask = BitMask32.allOn()):
+
+        """ Similar to screenshot(), this sets up a temporary cube map
+        Texture which it uses to take a series of six snapshots of the
+        current scene, one in each of the six cube map directions.
+        This requires rendering a new frame.
+
+        Unlike screenshot(), source may only be a GraphicsWindow,
+        GraphicsBuffer, or DisplayRegion; it may not be a Texture.
+
+        camera should be the node to which the cubemap cameras will be
+        parented.  The default is the camera associated with source,
+        if source is a DisplayRegion, or base.camera otherwise."""
+
+        if source == None:
+            source = base.win
+
+        if camera == None:
+            if hasattr(source, "getCamera"):
+                camera = source.getCamera()
+            if camera == None:
+                camera = base.camera
+
+        if hasattr(source, "getWindow"):
+            source = source.getWindow()
+
+        rig = NodePath(namePrefix)
+        buffer = source.makeCubeMap(namePrefix, size, 1, rig, cameraMask)
+        if buffer == None:
+            raise StandardError, "Could not make cube map."
+
+        # Now render a frame to fill up the texture.
+        rig.reparentTo(camera)
+        base.graphicsEngine.renderFrame()
+
+        tex = buffer.getTexture()
+        saved = self.screenshot(namePrefix = namePrefix,
+                                defaultFilename = defaultFilename,
+                                source = tex)
+
+        base.graphicsEngine.removeWindow(buffer)
+        rig.removeNode()
+
+        return saved
 
     def movie(self, namePrefix = 'movie', duration = 1.0, fps = 30,
-              format = 'rgb', sd = 4, source = None):
+              format = 'png', sd = 4, source = None):
         """
-        movie(namePrefix = 'movie', duration=1.0, fps=30, format='rgb', sd=4)
-
         Spawn a task to capture a movie using the screenshot function.
         - namePrefix will be used to form output file names (can include
           path information (e.g. '/i/beta/frames/myMovie')
         - duration is the length of the movie in seconds
         - fps is the frame rate of the resulting movie
-        - format specifies output file format (e.g. rgb, bmp)
+        - format specifies output file format (e.g. png, bmp)
         - sd specifies number of significant digits for frame count in the
-          output file name (e.g. if sd = 4, movie_0001.rgb)
+          output file name (e.g. if sd = 4, movie_0001.png)
         - source is the Window, Buffer, DisplayRegion, or Texture from which
           to save the resulting images.  The default is the main window.
         """
-        if source == None:
-            source = self.win
         globalClock.setMode(ClockObject.MNonRealTime)
         globalClock.setDt(1.0/float(fps))
         t = taskMgr.add(self._movieTask, namePrefix + '_task')
-        t.endT = globalClock.getFrameTime() + duration
         t.frameIndex = 1
+        t.numFrames = int(duration * fps)
         t.source = source
         t.outputString = namePrefix + '_%0' + `sd` + 'd.' + format
         t.uponDeath = lambda state: globalClock.setMode(ClockObject.MNormal)
 
     def _movieTask(self, state):
-        currT = globalClock.getFrameTime()
-        if currT >= state.endT:
+        frameName = state.outputString % state.frameIndex
+        self.notify.info("Capturing frame: " + frameName)
+        self.screenshot(namePrefix = frameName, defaultFilename = 0,
+                        source = state.source)
+                
+        state.frameIndex += 1
+        if state.frameIndex > state.numFrames:
             return Task.done
         else:
-            frameName = state.outputString % state.frameIndex
-            self.notify.info("Capturing frame: " + frameName)
-            if isinstance(state.source, Texture):
-                if state.source.getZSize() > 1:
-                    state.source.writePages(HashFilename(frameName))
-                else:
-                    state.source.write(Filename(frameName))
-            else:
-                state.source.saveScreenshot(Filename(frameName))
-                
-            state.frameIndex += 1
             return Task.cont
 
     def __windowEvent(self, win):