import pickle
from six.moves import range

from fuel.utils import do_not_pickle_attributes


@do_not_pickle_attributes("non_pickable", "bulky_attr")
class TestClass(object):
    def __init__(self):
        self.load()

    def load(self):
        self.bulky_attr = list(range(100))
        self.non_pickable = lambda x: x


def test_do_not_pickle_attributes():
    cl = TestClass()

    dump = pickle.dumps(cl)

    loaded = pickle.loads(dump)
    assert loaded.bulky_attr == list(range(100))
    assert loaded.non_pickable is not None

import pickle
from six.moves import range

from fuel.utils import do_not_pickle_attributes, expand_axis_label


@do_not_pickle_attributes("non_pickable", "bulky_attr")
class TestClass(object):
    def __init__(self):
        self.load()

    def load(self):
        self.bulky_attr = list(range(100))
        self.non_pickable = lambda x: x


def test_do_not_pickle_attributes():
    cl = TestClass()

    dump = pickle.dumps(cl)

    loaded = pickle.loads(dump)
    assert loaded.bulky_attr == list(range(100))
    assert loaded.non_pickable is not None


def test_expand_axis_label():
    assert expand_axis_label('b') == 'batch'
    assert expand_axis_label('c') == 'channel'
    assert expand_axis_label('t') == 'time'
    assert expand_axis_label('0') == 'axis_0'
    assert expand_axis_label('1') == 'axis_1'
    assert expand_axis_label('0b') == '0b'
    assert expand_axis_label('') == ''

Add unit test for expand_axis_label
Add unit test for expand_axis_label

  import pickle
  from six.moves import range
  
- from fuel.utils import do_not_pickle_attributes
+ from fuel.utils import do_not_pickle_attributes, expand_axis_label
  
  
  @do_not_pickle_attributes("non_pickable", "bulky_attr")
  class TestClass(object):
      def __init__(self):
          self.load()
  
      def load(self):
          self.bulky_attr = list(range(100))
          self.non_pickable = lambda x: x
  
  
  def test_do_not_pickle_attributes():
      cl = TestClass()
  
      dump = pickle.dumps(cl)
  
      loaded = pickle.loads(dump)
      assert loaded.bulky_attr == list(range(100))
      assert loaded.non_pickable is not None
  
+ 
+ def test_expand_axis_label():
+     assert expand_axis_label('b') == 'batch'
+     assert expand_axis_label('c') == 'channel'
+     assert expand_axis_label('t') == 'time'
+     assert expand_axis_label('0') == 'axis_0'
+     assert expand_axis_label('1') == 'axis_1'
+     assert expand_axis_label('0b') == '0b'
+     assert expand_axis_label('') == ''
+ 
Add unit test for expand_axis_label
## Code Before:
import pickle
from six.moves import range

from fuel.utils import do_not_pickle_attributes


@do_not_pickle_attributes("non_pickable", "bulky_attr")
class TestClass(object):
    def __init__(self):
        self.load()

    def load(self):
        self.bulky_attr = list(range(100))
        self.non_pickable = lambda x: x


def test_do_not_pickle_attributes():
    cl = TestClass()

    dump = pickle.dumps(cl)

    loaded = pickle.loads(dump)
    assert loaded.bulky_attr == list(range(100))
    assert loaded.non_pickable is not None

## Instruction:
Add unit test for expand_axis_label
## Code After:
import pickle
from six.moves import range

from fuel.utils import do_not_pickle_attributes, expand_axis_label


@do_not_pickle_attributes("non_pickable", "bulky_attr")
class TestClass(object):
    def __init__(self):
        self.load()

    def load(self):
        self.bulky_attr = list(range(100))
        self.non_pickable = lambda x: x


def test_do_not_pickle_attributes():
    cl = TestClass()

    dump = pickle.dumps(cl)

    loaded = pickle.loads(dump)
    assert loaded.bulky_attr == list(range(100))
    assert loaded.non_pickable is not None


def test_expand_axis_label():
    assert expand_axis_label('b') == 'batch'
    assert expand_axis_label('c') == 'channel'
    assert expand_axis_label('t') == 'time'
    assert expand_axis_label('0') == 'axis_0'
    assert expand_axis_label('1') == 'axis_1'
    assert expand_axis_label('0b') == '0b'
    assert expand_axis_label('') == ''
