python-pickle 0.1.0

data/lib/python/pickle/protocol5.rb

@@ -0,0 +1,218 @@
+require 'python/pickle/protocol4'
+require 'python/pickle/instructions/byte_array8'
+require 'python/pickle/instructions/next_buffer'
+require 'python/pickle/instructions/readonly_buffer'
+
+module Python
+  module Pickle
+    class Protocol5 < Protocol4
+
+      # Opcodes for Pickle protocol 5.
+      #
+      # @see https://peps.python.org/pep-0574/
+      OPCODES = Protocol4::OPCODES + Set[
+        150, # BYTEARRAY8
+        151, # NEXT_BUFFER
+        152, # READONLY_BUFFER
+      ]
+
+      #
+      # Reads an instruction from the pickle stream.
+      #
+      # @return [Instruction]
+      #   The decoded instruction.
+      #
+      # @raise [InvalidFormat]
+      #   The pickle stream could not be parsed.
+      #
+      def read_instruction
+        case (opcode = @io.getbyte)
+        #
+        # Protocol 0 instructions
+        #
+        when 40 # MARK
+          Instructions::MARK
+        when 46 # STOP
+          Instructions::STOP
+        when 48 # POP
+          Instructions::POP
+        when 49 # POP_MARK
+          Instructions::POP_MARK
+        when 50 # DUP
+          Instructions::DUP
+        when 70 # FLOAT
+          Instructions::Float.new(read_float)
+        when 73 # INT
+          Instructions::Int.new(read_int)
+        when 76 # LONG
+          Instructions::Long.new(read_long)
+        when 78 # NONE
+          Instructions::NONE
+        when 82 # REDUCE
+          Instructions::REDUCE
+        when 83 # STRING
+          Instructions::String.new(read_string)
+        when 86 # UNICODE
+          Instructions::String.new(read_unicode_string)
+        when 97 # APPEND
+          Instructions::APPEND
+        when 98 # BUILD
+          Instructions::BUILD
+        when 99 # GLOBAL
+          Instructions::Global.new(read_nl_string,read_nl_string)
+        when 100 # DICT
+          Instructions::DICT
+        when 103 # GET
+          Instructions::Get.new(read_int)
+        when 108 # LIST
+          Instructions::LIST
+        when 112 # PUT
+          Instructions::Put.new(read_int)
+        when 115 # SETITEM
+          Instructions::SETITEM
+        when 116 # TUPLE
+          Instructions::TUPLE
+        #
+        # Protocol 1 instructions
+        #
+        when 41  # EMPTY_TUPLE
+          Instructions::EMPTY_TUPLE
+        when 71  # BINFLOAT
+          Instructions::BinFloat.new(read_float64_be)
+        when 75  # BININT1
+          Instructions::BinInt1.new(read_uint8)
+        when 84  # BINSTRING
+          length = read_uint32_le
+          string = @io.read(length)
+
+          Instructions::BinString.new(length,string)
+        when 85  # SHORT_BINSTRING
+          length = read_uint8
+          string = @io.read(length)
+
+          Instructions::ShortBinString.new(length,string)
+        when 88  # BINUNICODE
+          length = read_uint32_le
+          string = @io.read(length).force_encoding(Encoding::UTF_8)
+
+          Instructions::BinUnicode.new(length,string)
+        when 93  # EMPTY_LIST
+          Instructions::EMPTY_LIST
+        when 101 # APPENDS
+          Instructions::APPENDS
+        when 104 # BINGET
+          Instructions::BinGet.new(read_uint8)
+        when 106 # LONG_BINGET
+          Instructions::LongBinGet.new(read_uint32_le)
+        when 113 # BINPUT
+          Instructions::BinPut.new(read_uint8)
+        when 117 # SETITEMS
+          Instructions::SETITEMS
+        when 125 # EMPTY_DICT
+          Instructions::EMPTY_DICT
+        #
+        # Protocol 2 instructions
+        #
+        when 128 # PROT
+          Instructions::Proto.new(read_uint8)
+        when 129 # NEWOBJ
+          Instructions::NEWOBJ
+        when 130 # EXT1
+          Instructions::Ext1.new(read_uint8)
+        when 131 # EXT2
+          Instructions::Ext2.new(read_uint16_le)
+        when 132 # EXT4
+          Instructions::Ext4.new(read_uint32_le)
+        when 133 # TUPLE1
+          Instructions::TUPLE1
+        when 134 # TUPLE2
+          Instructions::TUPLE2
+        when 135 # TUPLE3
+          Instructions::TUPLE3
+        when 136 # NEWTRUE
+          Instructions::NEWTRUE
+        when 137 # NEWFALSE
+          Instructions::NEWFALSE
+        when 138 # LONG1
+          length = read_uint8
+          long   = read_int_le(length)
+
+          Instructions::Long1.new(length,long)
+        when 139 # LONG4
+          length = read_uint32_le
+          long   = read_int_le(length)
+
+          Instructions::Long4.new(length,long)
+        #
+        # Protocol 3 instructions
+        #
+        when 66 # BINBYTES
+          length = read_uint32_le
+          bytes  = @io.read(length)
+
+          Instructions::BinBytes.new(length,bytes)
+        when 67 # SHORT_BINBYTES
+          length = read_uint8
+          bytes  = @io.read(length)
+
+          Instructions::ShortBinBytes.new(length,bytes)
+        #
+        # Protocol 4 instructions
+        #
+        when 140 # SHORT_BINUNICODE
+          length = read_uint8
+          string = read_utf8_string(length)
+
+          Instructions::ShortBinUnicode.new(length,string)
+        when 141 # BINUNICODE8
+          length = read_uint64_le
+          string = read_utf8_string(length)
+
+          Instructions::BinUnicode8.new(length,string)
+        when 142 # BINBYTES8
+          length = read_uint64_le
+          bytes  = @io.read(length)
+
+          Instructions::BinBytes8.new(length,bytes)
+        when 143 # EMPTY_SET
+          Instructions::EMPTY_SET
+        when 144 # ADDITEMS
+          Instructions::ADDITEMS
+        when 145 # FROZENSET
+          Instructions::FROZENSET
+        when 146 # NEWOBJ_EX
+          Instructions::NEWOBJ_EX
+        when 147 # STACK_GLOBAL
+          Instructions::STACK_GLOBAL
+        when 148 # MEMOIZE
+          Instructions::MEMOIZE
+        when 149 # FRAME
+          length = read_uint64_le
+
+          enter_frame(read_frame(length))
+
+          Instructions::Frame.new(length)
+        #
+        # Protocol 5 instructions.
+        #
+        when 150 # BYTEARRAY8
+          length = read_uint64_le
+          bytes  = @io.read(length)
+
+          Instructions::ByteArray8.new(length,bytes)
+        when 151 # NEXT_BUFFER
+          Instructions::NEXT_BUFFER
+        when 152 # READONLY_BUFFER
+          Instructions::READONLY_BUFFER
+        else
+          raise(InvalidFormat,"invalid opcode (#{opcode.inspect}) for protocol 5")
+        end
+      ensure
+        if @io.eof? && !@io_stack.empty?
+          leave_frame
+        end
+      end
+
+    end
+  end
+end

data/lib/python/pickle/py_class.rb

@@ -0,0 +1,75 @@
+require 'python/pickle/py_object'
+
+module Python
+  module Pickle
+    #
+    # Represents a Python class.
+    #
+    class PyClass
+
+      # The namespace the Python class is defined within.
+      #
+      # @return [String]
+      attr_reader :namespace
+
+      # The name of the Python class.
+      #
+      # @return [String]
+      attr_reader :name
+
+      #
+      # Initializes the Python class.
+      #
+      # @param [String, nil] namespace
+      #   The namespace of the Python class.
+      #
+      # @param [String] name
+      #   The name of the Python class.
+      #
+      # @api private
+      #
+      def initialize(namespace=nil,name)
+        @namespace = namespace
+        @name      = name
+      end
+
+      #
+      # Initializes a new Python object from the Python class.
+      #
+      # @param [Array] args
+      #   Additional `__init__` arguments.
+      #
+      # @param [Hash{Symbol => Object}] kwargs
+      #   Additional `__init__` keyword arguments.
+      #
+      # @api private
+      #
+      def new(*args,**kwargs)
+        PyObject.new(self,*args,**kwargs)
+      end
+
+      #
+      # Converts the Python class into a String.
+      #
+      # @return [String]
+      #
+      def to_s
+        if @namespace
+          "#{@namespace}.#{@name}"
+        else
+          @name.to_s
+        end
+      end
+
+      #
+      # Inspects the Python object.
+      #
+      # @return [String]
+      #
+      def inspect
+        "#<#{self.class}: #{self}>"
+      end
+
+    end
+  end
+end

data/lib/python/pickle/py_object.rb

@@ -0,0 +1,141 @@
+require 'python/pickle/py_class'
+
+module Python
+  module Pickle
+    #
+    # Represents a Python object.
+    #
+    class PyObject
+
+      # The Python class of the Python object.
+      #
+      # @return [PyClass]
+      attr_reader :py_class
+
+      # The arguments used to initialize the Python object.
+      #
+      # @return [Array]
+      attr_reader :init_args
+
+      # The keyword arguments used to initialize the Python object.
+      #
+      # @return [Array]
+      attr_reader :init_kwargs
+
+      # The populated attributes of the Python object.
+      #
+      # @return [Hash{String => Object}]
+      attr_reader :attributes
+
+      #
+      # Initializes the Python object.
+      #
+      # @param [PyClass] py_class
+      #   The Python class of the Python object.
+      #
+      # @param [Array] args
+      #   Additional arguments used to initialize the Python object.
+      #
+      # @param [Hash{Symbol => Object}] kwargs
+      #   Additional keyword arguments used to initialize the Python object.
+      #
+      # @api private
+      #
+      def initialize(py_class,*args,**kwargs)
+        @py_class    = py_class
+        @init_args   = args
+        @init_kwargs = kwargs
+
+        @attributes = {}
+      end
+
+      #
+      # Fetches the attribute of the Python object.
+      #
+      # @param [String] attribute
+      #   The attribute name.
+      #
+      # @return [Object]
+      #   The attributes value.
+      #
+      # @raise [ArgumentError]
+      #   The Python object does not have an attribute of the given name.
+      #
+      # @example
+      #   obj.getattr('x')
+      #   # => 2
+      #
+      def getattr(attribute)
+        @attributes.fetch(attribute) do
+          raise(ArgumentError,"Python object has no attribute #{attribute.inspect}: #{self.inspect}")
+        end
+      end
+
+      #
+      # Sets an attribute in the Python object.
+      #
+      # @param [String] name
+      #   The attribute name.
+      #
+      # @param [Object] value
+      #   The new value for the attribute.
+      #
+      # @example
+      #   obj.setattr('x',2)
+      #   # => 2
+      #
+      def setattr(name,value)
+        @attributes[name] = value
+      end
+
+      #
+      # Sets the state of the Python object.
+      #
+      # @api private
+      #
+      def __setstate__(new_attributes)
+        @attributes = new_attributes
+      end
+
+      #
+      # Converts the Python object to a Hash.
+      #
+      # @return [Hash]
+      #
+      def to_h
+        @attributes
+      end
+
+      protected
+
+      #
+      # Allows for direct access to attributes.
+      #
+      # @example
+      #   obj.x = 2
+      #   obj.x
+      #   # => 2
+      #
+      def method_missing(method_name,*arguments,&block)
+        if method_name.end_with?('=')
+          attr = method_name[0..-2]
+
+          if arguments.length == 1 && !block
+            return @attributes[attr] = arguments[0]
+          else
+            super(method_name,*arguments,&block)
+          end
+        else
+          attr = method_name.to_s
+
+          if @attributes.has_key?(attr) && arguments.empty? && !block
+            return @attributes[attr]
+          else
+            super(method_name,*arguments,&block)
+          end
+        end
+      end
+
+    end
+  end
+end

data/lib/python/pickle/tuple.rb

@@ -0,0 +1,19 @@
+module Python
+  module Pickle
+    #
+    # Represents a Python `tuple` object.
+    #
+    class Tuple < Array
+
+      #
+      # Inspects the tuple.
+      #
+      # @return [String]
+      #
+      def inspect
+        "#<#{self.class}: #{super}>"
+      end
+
+    end
+  end
+end

data/lib/python/pickle/version.rb

@@ -0,0 +1,6 @@
+module Python
+  module Pickle
+    # python-pickle version
+    VERSION = '0.1.0'
+  end
+end

data/lib/python/pickle.rb

@@ -0,0 +1,226 @@
+require 'python/pickle/protocol0'
+require 'python/pickle/protocol1'
+require 'python/pickle/protocol2'
+require 'python/pickle/protocol3'
+require 'python/pickle/protocol4'
+require 'python/pickle/protocol5'
+require 'python/pickle/exceptions'
+require 'python/pickle/deserializer'
+require 'stringio'
+
+module Python
+  #
+  # A modern Ruby implementation of the Python Pickle serialization format.
+  #
+  module Pickle
+    # Mapping of protocol versions to protocol parsers.
+    #
+    # @api private
+    PROTOCOL_VERSIONS = {
+      0 => Protocol0,
+      1 => Protocol1,
+      2 => Protocol2,
+      3 => Protocol3,
+      4 => Protocol4,
+      5 => Protocol5
+    }
+
+    # The default protocol version to use.
+    #
+    # @api public
+    DEFAULT_PROTCOL = 4
+
+    # The highest protocol version supported.
+    #
+    # @api public
+    HIGHEST_PROTOCOL = 5
+
+    #
+    # Parses a Python pickle stream.
+    #
+    # @param [String, IO] data
+    #   The Python pickle stream to parse.
+    #
+    # @param [Integer, nil] protocol
+    #   The explicit protocol version to use. If `nil` the protocol version will
+    #   be inferred by inspecting the first two bytes of the stream.
+    #
+    # @yield [instruction]
+    #   If a block is given, it will be passed each parsed Pickle instruction.
+    #
+    # @yieldparam [Instruction] instruction
+    #   A parsed Pickle instruction from the Pickle stream.
+    #
+    # @return [Array<Instruction>]
+    #   All parsed Pickle instructions from the Pickle stream.
+    #
+    # @api public
+    #
+    def self.parse(data, protocol: nil, &block)
+      io = case data 
+           when String then StringIO.new(data)
+           when IO     then data
+           else
+             raise(ArgumentError,"argument must be either an IO object or a String: #{io.inspect}")
+           end
+
+      if protocol
+        if (protocol < 0) || (protocol > HIGHEST_PROTOCOL)
+          raise(ArgumentError,"protocol must be between 0 or #{HIGHEST_PROTOCOL}, but was #{protocol.inspect}")
+        end
+      else
+        protocol = infer_protocol_version(io)
+      end
+
+      protocol_class = PROTOCOL_VERSIONS.fetch(protocol)
+      protocol       = protocol_class.new(io)
+
+      return protocol.read(&block)
+    end
+
+    #
+    # Deserializes the Python Pickle stream into a Ruby object.
+    #
+    # @param [String, IO] data
+    #   The Python pickle stream to parse.
+    #
+    # @param [Integer, nil] protocol
+    #   The explicit protocol version to use. If `nil` the protocol version will
+    #   be inferred by inspecting the first two bytes of the stream.
+    #
+    # @api public
+    #
+    def self.load(data,**kwargs)
+      deserializer = Deserializer.new(**kwargs)
+
+      parse(data) do |instruction|
+        status, object = deserializer.execute(instruction)
+
+        if status == :halt
+          return object
+        end
+      end
+
+      raise(DeserializationError,"failed to deserialize any object data from stream")
+    end
+
+    #
+    # Deserializes a Python Pickle file.
+    #
+    # @param [String] path
+    #   The path of the file.
+    #
+    # @return [Object]
+    #   The deserialized object.
+    #
+    def self.load_file(path,**kwargs)
+      load(File.open(path,'rb'),**kwargs)
+    end
+
+    #
+    # Serializes the Ruby object into Python Pickle data.
+    #
+    # @param [Object] object
+    #   The Ruby object to serialize.
+    #
+    # @param [IO] output
+    #   The option output to write the Pickle data to.
+    #
+    # @param [Integer] protocol
+    #   The desired Python Pickle protocol to use.
+    #
+    # @api public
+    #
+    def self.dump(object,output=nil, protocol: DEFAULT_PROTOCOL)
+      if (protocol < 0) || (protocol > HIGHEST_PROTOCOL)
+        raise(ArgumentError,"protocol must be between 0 or #{HIGHEST_PROTOCOL}, but was #{protocol.inspect}")
+      end
+    end
+
+    #
+    # Infers the protocol version from the IO stream.
+    #
+    # @param [IO] io
+    #   The IO stream to inspect.
+    #
+    # @return [Integer]
+    #   The inferred Python Pickle protocol version.
+    #
+    # @raise [InvalidFormat]
+    #   Could not determine the Pickle version from the first two bytes of the
+    #   IO stream.
+    #
+    # @api private
+    #
+    def self.infer_protocol_version(io)
+      opcode = io.getbyte
+
+      begin
+        case opcode
+        when 0x80 # PROTO (added in protocol 2)
+          version = io.getbyte
+          io.ungetbyte(version)
+          return version
+        when 48,  # POP (protocol 0)
+             50,  # DUP (protocol 0)
+             70,  # FLOAT (protocol 0)
+             83,  # STRING (protocol 0)
+             86,  # UNICODE (protocol 0)
+             100, # DICT (protocol 0)
+             103, # GET (protocol 0)
+             108, # LIST (protocol 0)
+             112  # PUT (protocol 0)
+          0
+        when 41,  # EMPTY_TUPLE (protocol 1)
+             71,  # BINFLOAT (protocol 1)
+             75,  # BININT1 (protocol 1)
+             84,  # BINSTRING (protocol 1)
+             85,  # SHORT_BINSTRING (protocol 1)
+             88,  # BINUNICODE (protocol 1)
+             93,  # EMPTY_LIST (protocol 1)
+             101, # APPENDS (protocol 1)
+             113, # BINPUT (protocol 1)
+             117, # SETITEMS (protocol 1)
+             125  # EMPTY_DICT (protocol 1)
+          1
+        when 46 # STOP
+          # if we've read all the way to the end of the stream and still cannot
+          # find any protocol 0 or protocol 1 specific opcodes, assume protocol 0
+          0
+        when 73, # INT  (identical in both protocol 0 and 1)
+             76  # LONG (identical in both protocol 0 and 1)
+          chars = io.gets
+
+          begin
+            infer_protocol_version(io)
+          ensure
+            chars.each_byte.reverse_each { |b| io.ungetbyte(b) }
+          end
+        when 40,  # MARK    (identical in both protocol 0 and 1)
+             78,  # NONE    (identical in both protocol 0 and 1)
+             82,  # REDUCE  (identical in both protocol 0 and 1)
+             97,  # APPEND  (identical in both protocol 0 and 1)
+             98,  # BUILD   (identical in both protocol 0 and 1)
+             115, # SETITEM (identical in both protocol 0 and 1)
+             116  # TUPLE   (identical in both protocol 0 and 1)
+          infer_protocol_version(io)
+        when 99 # GLOBAL
+          first_nl_string  = io.gets
+          second_nl_string = io.gets
+
+          begin
+            infer_protocol_version(io)
+          ensure
+            # push the read bytes back into the IO stream
+            second_nl_string.each_byte.reverse_each { |b| io.ungetbyte(b) }
+            first_nl_string.each_byte.reverse_each  { |b| io.ungetbyte(b) }
+          end
+        else
+          raise(InvalidFormat,"cannot infer protocol version from opcode (#{opcode.inspect}) at position #{io.pos}")
+        end
+      ensure
+        io.ungetbyte(opcode)
+      end
+    end
+  end
+end