import random

import numpy as np

import pytest

import torch

from docarray import Document, DocumentArray

from docarray.array.chunk import ChunkArray

from docarray.array.match import MatchArray

from docarray.score import NamedScore

@pytest.mark.parametrize('init_args', [None, dict(id=123), Document()])

@pytest.mark.parametrize('copy', [True, False])

def test_construct_doc(init_args, copy):

Document(init_args, copy)

def test_doc_hash_identical():

d1 = Document(text='hello')

d2 = Document(text='hello')

assert hash(d1) != hash(d2)

assert d1 != d2

d1.id = d2.id

assert hash(d1) == hash(d2)

assert d1 == d2

def test_doc_id_setter():

d = Document(

chunks=[Document(id='c0')],

id='d0',

)

assert d.id == d.chunks[0].parent_id == 'd0'

def test_random_seed_no_effect_on_id():

da = DocumentArray()

for _ in range(10):

random.seed(0)

np.random.seed(0)

torch.manual_seed(0)

torch.cuda.manual_seed_all(0)

da.append(

Document(

chunks=[Document()],

)

)

assert len(set(da[...][:, 'id'])) == 20

def test_doc_hash_complicate_content():

d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)

d2 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)

assert d1 == d2

assert d2 == d1

assert hash(d1) == hash(d2)

def test_doc_difference_complicate_content():

# Here we ensure != is symmetric therefore we put d1 != d2 and d2 != d1

# The __eq__ at DocumentData level is implemented in docarray/document/data.py

d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)

d2 = Document(text='hello', embedding=np.array([1, 2, 4]), id=1)

assert d1 != d2

assert d2 != d1

d1 = Document(text='hello', embedding=np.array([1, 2, 3, 5]), id=1)

d2 = Document(text='hello', embedding=np.array([1, 2, 4]), id=1)

assert d1 != d2

assert d2 != d1

d1 = Document(text='hello', id=1)

d2 = Document(text='hello', embedding=np.array([1, 2, 4]), id=1)

assert d1 != d2

assert d2 != d1

def test_pop_field():

d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)

assert d1.non_empty_fields == ('id', 'text', 'embedding')

d1.pop('text')

assert d1.non_empty_fields == ('id', 'embedding')

d1.pop('id', 'embedding')

assert d1.non_empty_fields == tuple()

d1.pop('foobar')

with pytest.raises(AttributeError):

assert d1.foobar

def test_clear_fields():

d1 = Document(text='hello', embedding=np.array([1, 2, 3]), id=1)

d1.clear()

assert d1.non_empty_fields == tuple()

def test_exclusive_content():

d = Document(text='hello')

assert d.content_type == 'text'

d.blob = b'123'

assert d.blob

assert not d.text

assert not d.tensor

assert d.content_type == 'blob'

d.tensor = [1, 2, 3]

assert d.tensor

assert not d.blob

assert not d.text

assert d.content_type == 'tensor'

d.text = 'hello'

assert d.text

assert not d.blob

assert not d.tensor

assert d.content_type == 'text'

def test_content_setter():

d = Document()

assert not d.content_type

d.content = 'hello'

assert d.content_type == 'text'

d.content = None

assert not d.content_type

def test_chunks_matches_setter():

d = Document(chunks=[Document()], matches=[Document(), Document()])

assert len(d.chunks) == 1

assert len(d.matches) == 2

assert isinstance(d.chunks, DocumentArray)

assert isinstance(d.chunks, ChunkArray)

assert isinstance(d.matches, DocumentArray)

assert isinstance(d.matches, MatchArray)

def test_empty_doc_chunks_matches():

assert isinstance(Document().chunks, DocumentArray)

assert isinstance(Document().matches, DocumentArray)

assert isinstance(Document().matches, MatchArray)

assert isinstance(Document().chunks, ChunkArray)

d = Document()

d.chunks.append(Document())

assert isinstance(d.chunks, ChunkArray)

d.chunks = [Document(), Document()]

assert isinstance(d.chunks, ChunkArray)

def test_chunk_match_increase_granularity():

d = Document()

d.chunks.append(Document())

assert d.chunks[0].granularity == 1

assert id(d.chunks.reference_doc) == id(d)

d.matches.append(Document())

assert d.matches[0].adjacency == 1

assert id(d.matches.reference_doc) == id(d)

d = d.chunks[0]

d.chunks.append(Document())

assert d.chunks[0].granularity == 2

assert id(d.chunks.reference_doc) == id(d)

d.matches.append(Document())

assert d.matches[0].adjacency == 1

assert id(d.matches.reference_doc) == id(d)

def test_offset():

d1 = Document(offset=1.0)

d2 = Document()

d2.offset = 1.0

assert d1.offset == d2.offset == 1.0

def test_exclusive_content_2():

d = Document(text='hello', blob=b'sda')

assert len(d.non_empty_fields) == 2

d.content = b'sda'

assert d.content == b'sda'

assert 'blob' in d.non_empty_fields

d = Document(content='hello')

assert d.content_type == 'text'

d = Document(content=b'hello')

assert d.content_type == 'blob'

d = Document(content=[1, 2, 3])

assert d.content_type == 'tensor'

def test_get_attr_values():

d = Document(

**{

'id': '123',

'text': 'document',

'feature1': 121,

'name': 'name',

'tags': {'id': 'identity', 'a': 'b', 'c': 'd', 'e': [0, 1, {'f': 'g'}]},

}

)

d.scores['metric'] = NamedScore(value=42)

required_keys = [

'id',

'text',

'tags__name',

'tags__feature1',

'scores__metric__value',

'tags__c',

'tags__id',

'tags__e__2__f',

]

res = d._get_attributes(*required_keys)

assert len(res) == len(required_keys)

assert res[required_keys.index('id')] == '123'

assert res[required_keys.index('tags__feature1')] == 121

assert res[required_keys.index('tags__name')] == 'name'

assert res[required_keys.index('text')] == 'document'

assert res[required_keys.index('tags__c')] == 'd'

assert res[required_keys.index('tags__id')] == 'identity'

assert res[required_keys.index('scores__metric__value')] == 42

assert res[required_keys.index('tags__e__2__f')] == 'g'

required_keys_2 = ['tags', 'text']

res2 = d._get_attributes(*required_keys_2)

assert len(res2) == 2

assert res2[required_keys_2.index('text')] == 'document'

assert res2[required_keys_2.index('tags')] == d.tags

d = Document({'id': '123', 'tags': {'outterkey': {'innerkey': 'real_value'}}})

required_keys_3 = ['tags__outterkey__innerkey']

res3 = d._get_attributes(*required_keys_3)

assert res3 == 'real_value'

d = Document(content=np.array([1, 2, 3]))

res4 = np.stack(d._get_attributes(*['tensor']))

np.testing.assert_equal(res4, np.array([1, 2, 3]))

def test_set_get_mime():

a = Document()

a.mime_type = 'jpg'

assert a.mime_type == 'image/jpeg'

b = Document()

b.mime_type = 'jpeg'

assert b.mime_type == 'image/jpeg'

c = Document()

c.mime_type = '.jpg'

assert c.mime_type == 'image/jpeg'

def test_doc_content():

d = Document()

assert d.content is None

d.text = 'abc'

assert d.content == 'abc'

c = np.random.random([10, 10])

d.tensor = c

np.testing.assert_equal(d.content, c)

d.blob = b'123'

assert d.blob == b'123'

def test_dict_constructor():

d1 = Document(

uri='https://jina.ai', mime_type='text/plain', granularity=1, adjacency=3

)

d2 = Document(

dict(uri='https://jina.ai', mime_type='text/plain', granularity=1, adjacency=3)

)

d3 = Document(

{

'uri': 'https://jina.ai',

'mime_type': 'text/plain',

'granularity': 1,

'adjacency': 3,

}

)

assert d1 != d2

d1.id = None

d2.id = None

d3.id = None

assert d1 == d2 == d3

def test_unknown_fields_behavior():

d = Document(hello='world')

assert d.tags == {'hello': 'world'}

d = Document(hello='world', unknown_fields_handler='drop')

assert d.tags == {}

with pytest.raises(AttributeError):

d = Document(hello='world', unknown_fields_handler='raise')

def test_content_setter_as_proxy():

d = Document(content='hello')

assert d.content == 'hello'

assert 'content' not in d.non_empty_fields

assert 'text' in d.non_empty_fields

d.content = [1, 2, 3]

assert 'tensor' in d.non_empty_fields

assert 'text' not in d.non_empty_fields