<reponame>ethan-asapp/flambe
import pytest
import tempfile
from collections import abc, OrderedDict
import os

import torch
import dill
import mock
from ruamel.yaml.compat import StringIO
from ruamel.yaml import YAML

from typing import Mapping, Any, Optional

# from flambe.compile import yaml
from flambe import Component, save_state_to_file, load_state_from_file, load, save
from flambe.compile import Registrable, yaml, make_component, Schema
from flambe.compile.serialization import _extract_prefix


FLAMBE_SOURCE_KEY = '_flambe_source'
FLAMBE_CLASS_KEY = '_flambe_class'
FLAMBE_CONFIG_KEY = '_flambe_config'
FLAMBE_DIRECTORIES_KEY = '_flambe_directories'
KEEP_VARS_KEY = 'keep_vars'
VERSION_KEY = '_flambe_version'


def list_files(startpath):
    for root, dirs, files in os.walk(startpath):
        level = root.replace(startpath, '').count(os.sep)
        indent = ' ' * 4 * (level)
        print('{}{}/'.format(indent, os.path.basename(root)))
        subindent = ' ' * 4 * (level + 1)
        for f in files:
            print('{}{}'.format(subindent, f))


def check_mapping_equivalence(x, y, exclude_config=False):
    for key in x.keys():
        if key == KEEP_VARS_KEY or key == 'version':
            continue
        if key == FLAMBE_CONFIG_KEY and exclude_config:
            continue
        assert key in y
        if isinstance(x[key], abc.Mapping):
            check_mapping_equivalence(x[key], y[key], exclude_config=exclude_config)
        elif isinstance(x[key], torch.Tensor):
            assert isinstance(y[key], torch.Tensor)
            torch.equal(x[key], y[key])
        else:
            assert x[key] == y[key]


EXAMPLE_TRAINER_CONFIG = """
!Trainer
train_sampler: !BaseSampler
val_sampler: !BaseSampler
dataset: !TabularDataset
  train: [['']]
model: !RNNEncoder
  input_size: 300
  rnn_type: lstm
  n_layers: 2
  hidden_size: 256
loss_fn: !torch.NLLLoss
metric_fn: !Accuracy
optimizer: !torch.Adam
  params: []
max_steps: 2
iter_per_step: 2
"""


@pytest.fixture
def make_classes_2():

    class A(Component):

        def __init__(self, akw1=0, akw2=None):
            self.akw1 = akw1
            self.akw2 = akw2

    class B(Component):

        def __init__(self, bkw1=0, bkw2='', bkw3=99):
            self.bkw1 = bkw1
            self.bkw2 = bkw2
            self.bkw3 = bkw3

    class C(Component):

        def __init__(self, one, two):
            self.one = one
            self.two = two

    return A, B


class Basic(Component):
    pass


class Composite(Component):
    def __init__(self):
        self.leaf = Basic()


class BasicStateful(Component):
    def __init__(self, x):
        self.x = x
        self.register_attrs('x')
        # self.b = Basic()


class BasicStatefulTwo(Component, torch.nn.Module):
    def __init__(self, y):
        super().__init__()
        self.y = y
        self.register_attrs('y')


class IntermediateTorch(Component, torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.leaf = Basic()

# TODO fix usage for x
class IntermediateStatefulTorch(Component, torch.nn.Module):
    def __init__(self, x):
        super().__init__()
        self.leaf = BasicStateful(x=x)
        self.linear = torch.nn.Linear(2, 2)


class IntermediateTorchOnly(torch.nn.Module):
    def __init__(self, component):
        super().__init__()
        self.child = component
        self.linear = torch.nn.Linear(2, 2)


class RootTorch(Component):
    def __init__(self, x):
        super().__init__()
        self.model = IntermediateStatefulTorch(x=x)
        # self.linear = torch.nn.Linear(2, 2)
        self.optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()))
        self.lr_scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, 0.01)


class ComposableTorchStateful(Component, torch.nn.Module):
    def __init__(self, a: Component, b: int, c: torch.nn.Module):
        super().__init__()
        self.child = a
        self.other_data = b
        self.linear = c
        self.register_attrs('other_data')


class ComposableTorchStatefulPrime(Component, torch.nn.Module):
    def __init__(self, a: Component, b: int, c: torch.nn.Linear):
        super().__init__()
        self.child = a
        self.other_data = b
        self.linear = c

    def _state(self, state_dict, prefix, local_metadata):
        state_dict[prefix + 'other_data'] = self.other_data
        return state_dict

    def _load_state(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        assert prefix + 'other_data' in state_dict
        self.other_data = state_dict[prefix + 'other_data']


class ComposableTorchStatefulTorchOnlyChild(Component, torch.nn.Module):
    def __init__(self, a: Component, b: int, c: Component):
        super().__init__()
        self.child = a
        self.other_data = b
        self.torch_only = IntermediateTorchOnly(c)
        self.register_attrs('other_data')


class ComposableContainer(Component):
    def __init__(self, item: Any):
        self.item = item


class Org(Component):
    def __init__(self, item: Any, extra):
        self.item = item
        self.torch_only_extra = extra


def create_factory(class_):
    def _factory(from_config):
        if from_config:
            config = f"!{class_.__name__} {{}}\n"
            obj = yaml.load(config)()
            return obj
        else:
            obj = class_()
        return obj
    return _factory


@pytest.fixture
def basic_object():
    return create_factory(Basic)


@pytest.fixture
def nested_object():
    return create_factory(Composite)


@pytest.fixture
def basic_object_with_state():
    def _factory(from_config, x=-1):
        if from_config:
            config = f"!BasicStateful\nx: {x}\n"
            obj = yaml.load(config)()
            return obj
        else:
            obj = BasicStateful(x=x)
        return obj
    return _factory


@pytest.fixture
def alternating_nn_module_with_state():
    def _factory(from_config, x=-1):
        if from_config:
            config = f"!RootTorch\nx: {x}\n"
            obj = yaml.load(config)()
            return obj
        else:
            obj = RootTorch(x=x)
        return obj
    return _factory


def schema_builder():
    config = """
!Basic
"""
    obj = yaml.load(config)
    return obj


def complex_builder(from_config, schema=False, x=-1):
    if from_config:
        config = """
!ComposableTorchStateful
a: !ComposableTorchStateful
  a: !ComposableTorchStateful
    a: !BasicStateful
      x: {}
    b: 2021
    c: !torch.Linear
      in_features: 2
      out_features: 2
  b: 2022
  c: !torch.Linear
    in_features: 2
    out_features: 2
b: 2023
c: !torch.Linear
  in_features: 2
  out_features: 2
"""
        config.format(x)
        obj = yaml.load(config)
        if not schema:
            obj = obj()
        return obj
    else:
        a1 = BasicStateful(x=x)
        b1 = 2021
        c1 = torch.nn.Linear(2, 2)
        a2 = ComposableTorchStateful(a1, b1, c1)
        b2 = 2022
        c2 = torch.nn.Linear(2, 2)
        a3 = ComposableTorchStateful(a2, b2, c2)
        b3 = 2023
        c3 = torch.nn.Linear(2, 2)
        obj = ComposableTorchStateful(a3, b3, c3)
        return obj


def complex_builder_nontorch_root(from_config, schema=False, x=-1):
    if from_config:
        config = """
!ComposableContainer
item:
  !ComposableTorchStatefulPrime
  a: !ComposableTorchStateful
    a: !ComposableTorchStateful
      a: !BasicStateful
        x: {}
      b: 2021
      c: !torch.Linear
        in_features: 2
        out_features: 2
    b: 2022
    c: !torch.Linear
      in_features: 2
      out_features: 2
  b: 2023
  c: !torch.Linear
    in_features: 2
    out_features: 2
"""
        config.format(x)
        obj = yaml.load(config)
        if not schema:
            obj = obj()
        return obj
    else:
        a1 = BasicStateful(x=x)
        b1 = 2021
        c1 = torch.nn.Linear(2, 2)
        a2 = ComposableTorchStateful(a1, b1, c1)
        b2 = 2022
        c2 = torch.nn.Linear(2, 2)
        a3 = ComposableTorchStateful(a2, b2, c2)
        b3 = 2023
        c3 = torch.nn.Linear(2, 2)
        item = ComposableTorchStateful(a3, b3, c3)
        obj = ComposableContainer(item)
        return obj


@pytest.fixture
def complex_multi_layered():
    return complex_builder


@pytest.fixture
def complex_multi_layered_nontorch_root():
    return complex_builder_nontorch_root


@pytest.fixture
def schema():
    return schema_builder

class TestHelpers:

    def test_extract_prefix(self):
        _extract_prefix


class TestState:

    def test_state_returns_not_None(self, basic_object):
        obj = basic_object(from_config=True)
        assert obj.get_state() is not None

    def test_state_metadata(self, basic_object):
        state = basic_object(from_config=True).get_state()
        assert hasattr(state, '_metadata')
        assert '' in state._metadata
        assert FLAMBE_DIRECTORIES_KEY in state._metadata
        assert FLAMBE_SOURCE_KEY in state._metadata['']
        assert VERSION_KEY in state._metadata['']
        assert state._metadata[''][FLAMBE_SOURCE_KEY] == "class Basic(Component):\n    pass\n"
        # assert state[FLAMBE_CONFIG_KEY] == ''
        assert '' in state._metadata[FLAMBE_DIRECTORIES_KEY] and len(state._metadata[FLAMBE_DIRECTORIES_KEY]) == 1
        assert state._metadata[''][VERSION_KEY] == '0.0.0'

    def test_state_config(self, basic_object):
        assert FLAMBE_CONFIG_KEY not in basic_object(from_config=False).get_state()._metadata['']
        obj = basic_object(from_config=True)
        state = obj.get_state()
        assert FLAMBE_CONFIG_KEY in state._metadata['']
        assert state._metadata[''][FLAMBE_CONFIG_KEY] == "!Basic {}\n"

    def test_state_nested_but_empty(self, nested_object):
        expected_state = {}
        expected_metadata = {'': {FLAMBE_SOURCE_KEY: "class Composite(Component):\n    def __init__(self):\n        self.leaf = Basic()\n", VERSION_KEY: "0.0.0", FLAMBE_CLASS_KEY: 'Composite'}, 'leaf': {FLAMBE_SOURCE_KEY: 'class Basic(Component):\n    pass\n', VERSION_KEY: '0.0.0', FLAMBE_CLASS_KEY: 'Basic'}, FLAMBE_DIRECTORIES_KEY: {'', 'leaf'}, KEEP_VARS_KEY: False}
        obj = nested_object(from_config=False)
        state = obj.get_state()
        assert state == expected_state
        check_mapping_equivalence(expected_metadata, state._metadata)
        check_mapping_equivalence(state._metadata, expected_metadata)

    def test_state_custom(self, basic_object_with_state):
        obj = basic_object_with_state(from_config=True)
        x = obj.x
        expected_state = {'x': x}
        assert obj.get_state() == expected_state

    # def test_state_custom_nested(nested_object_with_state):
    #     obj = nested_object_with_state()
    #     expected_state = {}
    #     assert obj.get_state() == expected_state
    #
    # def test_state_pytorch_empty(nn_modules):
    #     cls, cls_torch_first = nn_modules
    #     obj, obj_torch_first = cls(), cls_torch_first()
    #     expected_state = {}
    #     assert obj.get_state() == expected_state
    #     assert obj_torch_first.get_state() == expected_state
    #
    # def test_state_pytorch_nested_no_modules_no_parameters(nested_nn_module):
    #     obj = nested_nn_module()
    #     expected_state = {}
    #     assert obj.get_state() == expected_state
    #
    # def test_state_pytorch_alternating_nesting(alternating_nn_module):
    #     obj = alternating_nn_module()
    #     expected_state = {}
    #     assert obj.get_state() == expected_state

    def test_state_pytorch_alternating_nested_with_modules(self, alternating_nn_module_with_state):
        obj = alternating_nn_module_with_state(from_config=True, x=1)
        t1 = obj.model.linear.weight
        t2 = obj.model.linear.bias
        x = obj.model.leaf.x
        expected_state = {'model.leaf.x': x, 'model.linear.weight': t1, 'model.linear.bias': t2}
        root_source_code = dill.source.getsource(RootTorch)
        intermediate_source_code = dill.source.getsource(IntermediateStatefulTorch)
        leaf_source_code = dill.source.getsource(BasicStateful)
        expected_metadata = OrderedDict({FLAMBE_DIRECTORIES_KEY: set(['', 'model', 'model.leaf']), 'keep_vars': False, '': {VERSION_KEY: '0.0.0', FLAMBE_CLASS_KEY: 'RootTorch', FLAMBE_SOURCE_KEY: root_source_code, FLAMBE_CONFIG_KEY: "!RootTorch\nx: 1\n"},
                                                                                             'model': {VERSION_KEY: '0.0.0', FLAMBE_CLASS_KEY: 'IntermediateStatefulTorch', FLAMBE_SOURCE_KEY: intermediate_source_code, 'version': 1},  # TODO add config back: FLAMBE_CONFIG_KEY: "!IntermediateStatefulTorch {}\n"
                                                                                             'model.leaf': {VERSION_KEY: '0.0.0', FLAMBE_CLASS_KEY: 'BasicStateful', FLAMBE_SOURCE_KEY: leaf_source_code},  # TODO add config back: FLAMBE_CONFIG_KEY: "!BasicStateful {}\n"
                                                                                             'model.linear': {'version': 1}})
        state = obj.get_state()
        check_mapping_equivalence(state._metadata, expected_metadata)
        check_mapping_equivalence(expected_metadata, state._metadata)
        check_mapping_equivalence(state, expected_state)
        check_mapping_equivalence(expected_state, state)


class TestLoadState:

    def test_load_state_empty(self):
        pass

    def test_load_state_nested_empty(self):
        pass

    def test_load_state_custom_nested(self):
        pass

    def test_load_state_pytorch(self):
        pass

    def test_load_state_pytorch_alternating_nested(self):
        pass

    def test_state_complex_multilayered_nontorch_root(self, complex_multi_layered_nontorch_root):
        TORCH_TAG_PREFIX = "torch"
        exclude = ['torch.nn.quantized', 'torch.nn.qat']
        make_component(
            torch.nn.Module,
            TORCH_TAG_PREFIX,
            only_module='torch.nn',
            exclude=exclude
        )

        obj = complex_multi_layered_nontorch_root(from_config=True, x=1)
        t1 = obj.item.child.linear.weight.data
        state = obj.get_state()
        new_obj = complex_multi_layered_nontorch_root(from_config=True, x=2)
        new_obj.load_state(state)
        t2 = new_obj.item.child.linear.weight.data
        assert t1.equal(t2)
        check_mapping_equivalence(new_obj.get_state(), obj.get_state())
        check_mapping_equivalence(obj.get_state(), new_obj.get_state())

    def test_custom_attrs_load(self, complex_multi_layered):
        obj = complex_multi_layered(False)
        state = obj.state_dict()
        with pytest.raises(RuntimeError) as excinfo:
            obj.load_state_dict(state, strict=True)
        assert "Unexpected key(s)" in str(excinfo.value)
        obj.load_state_dict(state, strict=False)


class TestClassSave:

    def test_class_save(self):
        pass


class TestClassLoad:

    def test_class_load(self):
        pass


class TestModuleSave:

    def test_save_single_object(self, basic_object):
        pass

    def test_save_nested_object(self, nested_object):
        pass

    def test_save_pytorch_nested_alternating(self, alternating_nn_module_with_state):
        pass


class TestModuleLoad:

    def test_load_directory_single_file(self, basic_object):
        pass

    def test_load_nested_directory(self, nested_object):
        pass

    def test_load_pytorch_alternating(self, alternating_nn_module_with_state):
        pass


class TestSerializationIntegration:

    def test_state_and_load_roundtrip_single_object(self, basic_object):
        old_obj = basic_object(from_config=True)
        state = old_obj.get_state()
        new_obj = basic_object(from_config=False)
        new_obj.load_state(state, strict=False)
        assert old_obj.get_state() == new_obj.get_state()

    # def test_state_and_load_roundtrip_nested_object(self):
    #     pass

    def test_state_and_load_roundtrip_pytorch_alternating(self, alternating_nn_module_with_state):
        old_obj = alternating_nn_module_with_state(from_config=True, x=1)
        state = old_obj.get_state()
        new_obj = alternating_nn_module_with_state(from_config=False, x=2)
        new_obj.load_state(state, strict=False)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=True)

    # def test_class_save_and_load_roundtrip():
    #     pass
    #
    # def test_class_save_and_load_roundtrip_nested():
    #     pass
    #
    # def test_class_save_and_load_roundtrip_pytorch():
    #     pass

    def test_save_to_file_and_load_from_file_roundtrip(self, basic_object):
        old_obj = basic_object(from_config=True)
        state = old_obj.get_state()
        with tempfile.TemporaryDirectory() as path:
            save_state_to_file(state, path)
            state = load_state_from_file(path)
        new_obj = basic_object(from_config=False)
        new_obj.load_state(state, strict=False)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=True)

    def test_save_to_file_and_load_from_file_roundtrip_pytorch(self, alternating_nn_module_with_state):
        old_obj = alternating_nn_module_with_state(from_config=False, x=1)
        state = old_obj.get_state()
        with tempfile.TemporaryDirectory() as path:
            save_state_to_file(state, path)
            state = load_state_from_file(path)
        new_obj = alternating_nn_module_with_state(from_config=False, x=2)
        new_obj.load_state(state, strict=False)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)

    def test_save_to_file_and_load_from_file_roundtrip_complex(self, complex_multi_layered):
        TORCH_TAG_PREFIX = "torch"
        exclude = ['torch.nn.quantized', 'torch.nn.qat']
        make_component(
            torch.nn.Module,
            TORCH_TAG_PREFIX,
            only_module='torch.nn',
            exclude=exclude
        )
        old_obj = complex_multi_layered(from_config=True, x=1)
        # Test that the current state is actually saved, for a
        # Component-only child of torch objects
        old_obj.child.child.child.x = 24
        state = old_obj.get_state()
        with tempfile.TemporaryDirectory() as path:
            save_state_to_file(state, path)
            list_files(path)
            state_loaded = load_state_from_file(path)
            check_mapping_equivalence(state, state_loaded)
            # assert False
        new_obj = complex_multi_layered(from_config=True, x=2)
        new_obj.load_state(state_loaded, strict=False)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)

    @pytest.mark.parametrize("pickle_only", [True, False])
    @pytest.mark.parametrize("compress_save_file", [True, False])
    def test_save_to_file_and_load_from_file_roundtrip_complex_nontorch_root(self,
            complex_multi_layered_nontorch_root, pickle_only, compress_save_file):
        TORCH_TAG_PREFIX = "torch"
        exclude = ['torch.nn.quantized', 'torch.nn.qat']
        make_component(
            torch.nn.Module,
            TORCH_TAG_PREFIX,
            only_module='torch.nn',
            exclude=exclude
        )
        old_obj = complex_multi_layered_nontorch_root(from_config=True, x=1)
        state = old_obj.get_state()
        with tempfile.TemporaryDirectory() as root_path:
            path = os.path.join(root_path, 'savefile.flambe')
            save_state_to_file(state, path, compress_save_file, pickle_only)
            list_files(path)
            if pickle_only:
                path += '.pkl'
            if compress_save_file:
                path += '.tar.gz'
            state_loaded = load_state_from_file(path)
            check_mapping_equivalence(state, state_loaded)
            check_mapping_equivalence(state._metadata, state_loaded._metadata)
        new_obj = complex_multi_layered_nontorch_root(from_config=True, x=2)
        int_state = new_obj.get_state()
        new_obj.load_state(state_loaded, strict=False)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata)
        check_mapping_equivalence(int_state._metadata, state_loaded._metadata)

    @pytest.mark.parametrize("pickle_only", [True, False])
    @pytest.mark.parametrize("compress_save_file", [True, False])
    def test_module_save_and_load_roundtrip(self, basic_object, pickle_only, compress_save_file):
        old_obj = basic_object(from_config=True)
        with tempfile.TemporaryDirectory() as root_path:
            path = os.path.join(root_path, 'savefile.flambe')
            save(old_obj, path, compress_save_file, pickle_only)
            if pickle_only:
                path += '.pkl'
            if compress_save_file:
                path += '.tar.gz'
            new_obj = load(path)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)

    @pytest.mark.parametrize("pickle_only", [True, False])
    @pytest.mark.parametrize("compress_save_file", [True, False])
    def test_module_save_and_load_roundtrip_pytorch(self,
                                                    alternating_nn_module_with_state,
                                                    pickle_only,
                                                    compress_save_file):
        old_obj = alternating_nn_module_with_state(from_config=True, x=1)
        with tempfile.TemporaryDirectory() as root_path:
            path = os.path.join(root_path, 'savefile.flambe')
            save(old_obj, path, compress_save_file, pickle_only)
            if pickle_only:
                path += '.pkl'
            if compress_save_file:
                path += '.tar.gz'
            new_obj = load(path)
        old_state = old_obj.get_state()
        new_state = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)


    def test_module_save_and_load_roundtrip_pytorch_only_bridge(self):
        a = BasicStateful.compile(x=3)
        b = 100
        c = BasicStatefulTwo.compile(y=0)
        item = ComposableTorchStatefulTorchOnlyChild.compile(a=a, b=b, c=c)
        extra = torch.nn.Linear(2, 2)
        old_obj = Org.compile(item=item, extra=None)
        # x for a2 should be different from instance a
        a2 = BasicStateful.compile(x=4)
        b2 = 101
        # y for c2 should be different from instance c
        c2 = BasicStatefulTwo.compile(y=1)
        item2 = ComposableTorchStatefulTorchOnlyChild.compile(a=a2, b=b2, c=c2)
        extra2 = torch.nn.Linear(2, 2)
        new_obj = Org.compile(item=item2, extra=None)
        with tempfile.TemporaryDirectory() as root_path:
            path = os.path.join(root_path, 'asavefile2.flambe')
            old_state = old_obj.get_state()
            save_state_to_file(old_state, path)
            new_state = load_state_from_file(path)
            new_obj.load_state(new_state)
            # save(old_obj, path)
            # new_obj = load(path)
        old_state_get = old_obj.get_state()
        new_state_get = new_obj.get_state()
        check_mapping_equivalence(new_state, old_state)
        check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)
        check_mapping_equivalence(new_state_get, old_state_get)
        check_mapping_equivalence(old_state_get._metadata, new_state_get._metadata, exclude_config=True)

    # def test_module_save_and_load_example_encoder(self):
    #     TORCH_TAG_PREFIX = "torch"
    #     make_component(torch.nn.Module, TORCH_TAG_PREFIX, only_module='torch.nn')
    #     make_component(torch.optim.Optimizer, TORCH_TAG_PREFIX, only_module='torch.optim')
    #     trainer = yaml.load(EXAMPLE_TRAINER_CONFIG)()
    #     with tempfile.TemporaryDirectory() as path:
    #         save(trainer, path)
    #         loaded_trainer = load(path)
    #     old_state = trainer.get_state()
    #     new_state = loaded_trainer.get_state()
    #     check_mapping_equivalence(new_state, old_state)
    #     check_mapping_equivalence(old_state._metadata, new_state._metadata, exclude_config=False)

    def test_module_save_and_load_single_instance_appears_twice(self, make_classes_2):
        txt = """
!C
one: !A
  akw2: &theb !B
    bkw2: test
    bkw1: 1
  akw1: 8
two: !A
  akw1: 8
  # Comment Here
  akw2: *theb
"""
        c = yaml.load(txt)()
        c.one.akw2.bkw1 = 6
        assert c.one.akw2 is c.two.akw2
        assert c.one.akw2.bkw1 == c.two.akw2.bkw1
        with tempfile.TemporaryDirectory() as path:
            save(c, path)
            state = load_state_from_file(path)
            loaded_c = load(path)
        assert loaded_c.one.akw2 is loaded_c.two.akw2
        assert loaded_c.one.akw2.bkw1 == loaded_c.two.akw2.bkw1


class TestSerializationExtensions:

    EXTENSIONS = {
        "ext1": "my_extension_1",
        "ext2": "my_extension_2",
        "ext3": "my_extension_3",
    }

    @pytest.mark.parametrize("pickle_only", [True, False])
    @pytest.mark.parametrize("compress_save_file", [True, False])
    @mock.patch('flambe.compile.serialization.is_installed_module')
    @mock.patch('flambe.compile.serialization.import_modules')
    @mock.patch('flambe.compile.component.Schema.add_extensions_metadata')
    def test_save_to_file_and_load_from_file_with_extensions(
            self, mock_add_extensions,
            mock_import_module, mock_installed_module,
            compress_save_file, pickle_only, schema):
        """Test that extensions are saved to the output config.yaml
        and they are also added when loading back the object."""

        mock_installed_module.return_value = True

        schema_obj = schema()

        # Add extensions manually because if we use add_extensions_metadata
        # then no extensions will be added as the schema doesn't container_folder
        # any prefix.
        schema_obj._extensions = TestSerializationExtensions.EXTENSIONS

        obj = schema_obj()
        state = obj.get_state()

        with tempfile.TemporaryDirectory() as root_path:
            path = os.path.join(root_path, 'savefile.flambe')
            save_state_to_file(state, path, compress_save_file, pickle_only)

            list_files(path)
            if pickle_only:
                path += '.pkl'
            if compress_save_file:
                path += '.tar.gz'
            state_loaded = load_state_from_file(path)

            check_mapping_equivalence(state, state_loaded)
            check_mapping_equivalence(state._metadata, state_loaded._metadata)

            _ = Basic.load_from_path(path)
            mock_add_extensions.assert_called_once_with(TestSerializationExtensions.EXTENSIONS)


    def test_add_extensions_metadata(self, schema):
        """Test that add_extensions_metadata doesn't add extensions that are not used"""

        schema_obj = schema()
        assert schema_obj._extensions == {}

        schema_obj.add_extensions_metadata(TestSerializationExtensions.EXTENSIONS)
        assert schema_obj._extensions == {}


    def test_add_extensions_metadata_2(self):
        """Test that add_extensions_metadata doesn't add extensions that are not used.

        In this case we will use a config containing torch, but we will make_component
        on torch so that it can be compiled. After that, we add_extensions_metadata with
        torch, which is a valid extensions for the config (redundant, but valid).

        """
        TORCH_TAG_PREFIX = "torch"
        exclude = ['torch.nn.quantized', 'torch.nn.qat']
        make_component(
            torch.nn.Module,
            TORCH_TAG_PREFIX,
            only_module='torch.nn',
            exclude=exclude
        )

        config = """
        !torch.Linear
          in_features: 2
          out_features: 2
        """

        schema = yaml.load(config)
        schema.add_extensions_metadata({"torch": "torch"})
        assert schema._extensions == {"torch": "torch"}

        mixed_ext = TestSerializationExtensions.EXTENSIONS.copy()
        mixed_ext.update({"torch": "torch"})
        schema.add_extensions_metadata(mixed_ext)
        assert schema._extensions == {"torch": "torch"}


    def test_add_extensions_metadata_3(self, complex_multi_layered_nontorch_root):
        """Test that add_extensions_metadata doesn't add extensions that are not used

        In this case we will use a config containing torch, but we will make_component
        on torch so that it can be compiled. After that, we add_extensions_metadata with
        torch, which is a valid extensions for the config (redundant, but valid).

        """
        TORCH_TAG_PREFIX = "torch"
        exclude = ['torch.nn.quantized', 'torch.nn.qat']
        make_component(
            torch.nn.Module,
            TORCH_TAG_PREFIX,
            only_module='torch.nn',
            exclude=exclude
        )

        schema = complex_multi_layered_nontorch_root(from_config=True, schema=True)
        schema.add_extensions_metadata({"torch": "torch"})

        # This method asserts recursively that torch is added to extensions when
        # there is a subcomponent that uses torch.
        # It returns if at least one component with torch was found, that should
        # always happen based on the complex_multi_layered_nontorch_root.
        def helper(data):
            found = False
            if isinstance(data, Schema):
                if data.component_subclass.__module__.startswith("torch."):
                    found = True
                    assert data._extensions == {"torch": "torch"}

                for val in data.keywords.values():
                    f = helper(val)
                    if f:
                        found = f

            elif isinstance(data, Mapping):
                for val in data.values():
                    f = helper(val)
                    if f:
                        found = f
            return found

        assert helper(schema)
