from pymongo import *
from bson import *
from bson.codec_options import *


class UnecodableType(object):
    def __init__(self, num):
        self.num = num


class DecodeInt64AsUnecodableType(TypeCodecBase):
    @property
    def bson_type(self):
        return Int64

    def transform_bson(self, value):
        return UnecodableType(value)

    # Encoding this type is not supported.
    # @property
    # def python_type(self):
    #     return MyCustomInt64
    # def transform_python(self, value):
    #     return Int64(value.num)


client = MongoClient()
coll = client.test.test
coll.drop()
coll.insert({'_id': i} for i in range(5))

cursor = coll.find(batch_size=2)
cursor.next()
print('cursor.cursor_id: %r' % cursor.cursor_id)
print(list(cursor))

type_registry = TypeRegistry((DecodeInt64AsUnecodableType(),))
codec_options = CodecOptions(type_registry=type_registry)
coll_with_decoder = coll.with_options(codec_options=codec_options)
cursor = coll_with_decoder.find(batch_size=2)
cursor.next()
print('cursor.cursor_id: %r' % cursor.cursor_id)
print(list(cursor))

cursor.cursor_id: 4026529133242359067L
[{u'_id': 1}, {u'_id': 2}, {u'_id': 3}, {u'_id': 4}]
cursor.cursor_id: <__main__.UnecodableType object at 0x10baf5ad0>
Traceback (most recent call last):
  File "type_codec.py", line 44, in <module>
    print(list(cursor))
  File "/Users/shane/git/mongo-python-driver/pymongo/cursor.py", line 1189, in next
    if len(self.__data) or self._refresh():
  File "/Users/shane/git/mongo-python-driver/pymongo/cursor.py", line 1126, in _refresh
    self.__send_message(g)
  File "/Users/shane/git/mongo-python-driver/pymongo/cursor.py", line 931, in __send_message
    operation, exhaust=self.__exhaust, address=self.__address)
  File "/Users/shane/git/mongo-python-driver/pymongo/mongo_client.py", line 1211, in _send_message_with_response
    exhaust)
  File "/Users/shane/git/mongo-python-driver/pymongo/server.py", line 94, in send_message_with_response
    set_slave_okay, sock_info, use_find_cmd)
  File "/Users/shane/git/mongo-python-driver/pymongo/message.py", line 399, in get_message
    ctx=sock_info.compression_context)
  File "/Users/shane/git/mongo-python-driver/pymongo/message.py", line 678, in _op_msg
    flags, command, identifier, docs, check_keys, opts)
bson.errors.InvalidDocument: Cannot encode object: <__main__.UnecodableType object at 0x10baf5ad0>