<s> import logging <EOL> import json <EOL> import os <EOL> import subprocess <EOL> import tempfile <EOL> import os . path <EOL> import distutils . spawn , distutils . version <EOL> import numpy as np <EOL> from six import StringIO <EOL> import six <EOL> import six . moves . urllib as urlparse <EOL> from gym import error <EOL> logger = logging . getLogger ( __name__ ) <EOL> def touch ( path ) : <EOL> open ( path , '<STR_LIT:a>' ) . close ( ) <EOL> class VideoRecorder ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , env , path = None , metadata = None , enabled = True , base_path = None ) : <EOL> modes = env . metadata . get ( '<STR_LIT>' , [ ] ) <EOL> self . ansi_mode = False <EOL> if '<STR_LIT>' not in modes : <EOL> if '<STR_LIT>' in modes : <EOL> self . ansi_mode = True <EOL> else : <EOL> logger . info ( '<STR_LIT>' . format ( env ) ) <EOL> enabled = False <EOL> if path is not None and base_path is not None : <EOL> raise error . Error ( "<STR_LIT>" ) <EOL> self . enabled = enabled <EOL> self . last_frame = None <EOL> if not self . enabled : <EOL> return <EOL> self . env = env <EOL> required_ext = '<STR_LIT>' if self . ansi_mode else '<STR_LIT>' <EOL> if path is None : <EOL> if base_path is not None : <EOL> path = base_path + required_ext <EOL> else : <EOL> with tempfile . NamedTemporaryFile ( suffix = required_ext , delete = False ) as f : <EOL> path = f . name <EOL> self . path = path <EOL> path_base , actual_ext = os . path . splitext ( self . path ) <EOL> if actual_ext != required_ext : <EOL> hint = "<STR_LIT>" if self . ansi_mode else '<STR_LIT>' <EOL> raise error . Error ( "<STR_LIT>" . format ( self . path , required_ext , hint ) ) <EOL> touch ( path ) <EOL> self . frames_per_sec = env . metadata . get ( '<STR_LIT>' , <NUM_LIT:30> ) <EOL> self . encoder = None <EOL> self . broken = False <EOL> self . metadata = metadata or { } <EOL> self . metadata [ '<STR_LIT>' ] = '<STR_LIT>' if self . ansi_mode else '<STR_LIT>' <EOL> self . metadata_path = '<STR_LIT>' . format ( path_base ) <EOL> self . write_metadata ( ) <EOL> logger . info ( '<STR_LIT>' , self . path ) <EOL> self . empty = True <EOL> @ property <EOL> def functional ( self ) : <EOL> return self . enabled and not self . broken <EOL> def capture_frame ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . functional : return <EOL> logger . debug ( '<STR_LIT>' , self . path ) <EOL> render_mode = '<STR_LIT>' if self . ansi_mode else '<STR_LIT>' <EOL> frame = self . env . render ( mode = render_mode ) <EOL> if frame is None : <EOL> logger . warn ( '<STR_LIT>' , self . path , self . metadata_path ) <EOL> self . broken = True <EOL> else : <EOL> self . last_frame = frame <EOL> if self . ansi_mode : <EOL> self . _encode_ansi_frame ( frame ) <EOL> else : <EOL> self . _encode_image_frame ( frame ) <EOL> def close ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . enabled : <EOL> return <EOL> if self . encoder : <EOL> logger . debug ( '<STR_LIT>' , self . path ) <EOL> self . encoder . close ( ) <EOL> self . encoder = None <EOL> else : <EOL> os . remove ( self . path ) <EOL> if self . metadata is None : <EOL> self . metadata = { } <EOL> self . metadata [ '<STR_LIT>' ] = True <EOL> if self . broken : <EOL> logger . info ( '<STR_LIT>' , self . path , self . metadata_path ) <EOL> if os . path . exists ( self . path ) : <EOL> os . remove ( self . path ) <EOL> if self . metadata is None : <EOL> self . metadata = { } <EOL> self . metadata [ '<STR_LIT>' ] = True <EOL> self . write_metadata ( ) <EOL> def write_metadata ( self ) : <EOL> with open ( self . metadata_path , '<STR_LIT:w>' ) as f : <EOL> json . dump ( self . metadata , f ) <EOL> def _encode_ansi_frame ( self , frame ) : <EOL> if not self . encoder : <EOL> self . encoder = TextEncoder ( self . path , self . frames_per_sec ) <EOL> self . metadata [ '<STR_LIT>' ] = self . encoder . version_info <EOL> self . encoder . capture_frame ( frame ) <EOL> self . empty = False <EOL> def _encode_image_frame ( self , frame ) : <EOL> if not self . encoder : <EOL> self . encoder = ImageEncoder ( self . path , frame . shape , self . frames_per_sec ) <EOL> self . metadata [ '<STR_LIT>' ] = self . encoder . version_info <EOL> try : <EOL> self . encoder . capture_frame ( frame ) <EOL> except error . InvalidFrame as e : <EOL> logger . warn ( '<STR_LIT>' , e ) <EOL> self . broken = True <EOL> else : <EOL> self . empty = False <EOL> class TextEncoder ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , output_path , frames_per_sec ) : <EOL> self . output_path = output_path <EOL> self . frames_per_sec = frames_per_sec <EOL> self . frames = [ ] <EOL> def capture_frame ( self , frame ) : <EOL> string = None <EOL> if isinstance ( frame , str ) : <EOL> string = frame <EOL> elif isinstance ( frame , StringIO ) : <EOL> string = frame . getvalue ( ) <EOL> else : <EOL> raise error . InvalidFrame ( '<STR_LIT>' . format ( type ( frame ) , frame ) ) <EOL> frame_bytes = string . encode ( '<STR_LIT:utf-8>' ) <EOL> if frame_bytes [ - <NUM_LIT:1> : ] != six . b ( '<STR_LIT:\n>' ) : <EOL> raise error . InvalidFrame ( '<STR_LIT>' . format ( string ) ) <EOL> if six . b ( '<STR_LIT:\r>' ) in frame_bytes : <EOL> raise error . InvalidFrame ( '<STR_LIT>' . format ( string ) ) <EOL> self . frames . append ( frame_bytes ) <EOL> def close ( self ) : <EOL> frame_duration = <NUM_LIT> <EOL> clear_code = six . b ( "<STR_LIT>" % ( <NUM_LIT> ) ) <EOL> events = [ ( frame_duration , ( clear_code + frame . replace ( six . b ( '<STR_LIT:\n>' ) , six . b ( '<STR_LIT:\r\n>' ) ) ) . decode ( '<STR_LIT:utf-8>' ) ) for frame in self . frames ] <EOL> height = max ( [ frame . count ( six . b ( '<STR_LIT:\n>' ) ) for frame in self . frames ] ) + <NUM_LIT:1> <EOL> width = max ( [ max ( [ len ( line ) for line in frame . split ( six . b ( '<STR_LIT:\n>' ) ) ] ) for frame in self . frames ] ) + <NUM_LIT:2> <EOL> data = { <EOL> "<STR_LIT:version>" : <NUM_LIT:1> , <EOL> "<STR_LIT:width>" : width , <EOL> "<STR_LIT>" : height , <EOL> "<STR_LIT>" : len ( self . frames ) * frame_duration , <EOL> "<STR_LIT>" : "<STR_LIT:->" , <EOL> "<STR_LIT:title>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : { } , <EOL> "<STR_LIT>" : events , <EOL> } <EOL> with open ( self . output_path , '<STR_LIT:w>' ) as f : <EOL> json . dump ( data , f ) <EOL> @ property <EOL> def version_info ( self ) : <EOL> return { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:version>' : <NUM_LIT:1> } <EOL> class ImageEncoder ( object ) : <EOL> def __init__ ( self , output_path , frame_shape , frames_per_sec ) : <EOL> self . proc = None <EOL> self . output_path = output_path <EOL> h , w , pixfmt = frame_shape <EOL> if pixfmt != <NUM_LIT:3> and pixfmt != <NUM_LIT:4> : <EOL> raise error . InvalidFrame ( "<STR_LIT>" . format ( frame_shape ) ) <EOL> self . wh = ( w , h ) <EOL> self . includes_alpha = ( pixfmt == <NUM_LIT:4> ) <EOL> self . frame_shape = frame_shape <EOL> self . frames_per_sec = frames_per_sec <EOL> if distutils . spawn . find_executable ( '<STR_LIT>' ) is not None : <EOL> self . backend = '<STR_LIT>' <EOL> elif distutils . spawn . find_executable ( '<STR_LIT>' ) is not None : <EOL> self . backend = '<STR_LIT>' <EOL> else : <EOL> raise error . DependencyNotInstalled ( """<STR_LIT>""" ) <EOL> self . start ( ) <EOL> @ property <EOL> def version_info ( self ) : <EOL> return { '<STR_LIT>' : self . backend , '<STR_LIT:version>' : subprocess . check_output ( [ self . backend , '<STR_LIT>' ] ) , '<STR_LIT>' : self . cmdline } <EOL> def start ( self ) : <EOL> self . cmdline = ( self . backend , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT:error>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' % self . frames_per_sec , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( * self . wh ) , <EOL> '<STR_LIT>' , ( '<STR_LIT>' if self . includes_alpha else '<STR_LIT>' ) , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> self . output_path <EOL> ) <EOL> logger . debug ( '<STR_LIT>' , '<STR_LIT:U+0020>' . join ( self . cmdline ) ) <EOL> self . proc = subprocess . Popen ( self . cmdline , stdin = subprocess . PIPE ) <EOL> def capture_frame ( self , frame ) : <EOL> if not isinstance ( frame , ( np . ndarray , np . generic ) ) : <EOL> raise error . InvalidFrame ( '<STR_LIT>' . format ( type ( frame ) , frame ) ) <EOL> if frame . shape != self . frame_shape : <EOL> raise error . InvalidFrame ( "<STR_LIT>" . format ( frame_shape , self . frame_shape ) ) <EOL> if frame . dtype != np . uint8 : <EOL> raise error . InvalidFrame ( "<STR_LIT>" . format ( frame . dtype ) ) <EOL> if distutils . version . StrictVersion ( np . __version__ ) >= distutils . version . StrictVersion ( '<STR_LIT>' ) : <EOL> self . proc . stdin . write ( frame . tobytes ( ) ) <EOL> else : <EOL> self . proc . stdin . write ( frame . tostring ( ) ) <EOL> def close ( self ) : <EOL> self . proc . stdin . close ( ) <EOL> ret = self . proc . wait ( ) <EOL> if ret != <NUM_LIT:0> : <EOL> logger . error ( "<STR_LIT>" . format ( ret ) ) </s>
