utils/load_save.py

import yaml
import glob
import pickle
import os
import socket


def get_file_path(file_path):
    hostname = socket.gethostname()
    if hostname == 'fib':
        file_path = os.path.join('/home/mas/zhengyu/workspace/flp', file_path)
    elif hostname == 'rl2':
        file_path = os.path.join('/home/zhengyu/workspace/flp', file_path)
    elif hostname == 'rl3':
        file_path = os.path.join('/home/zhengyu/workspace/flp', file_path)
    elif hostname == 'rl4':
        file_path = os.path.join('/data2/suhongyuan/flp', file_path)
    elif hostname == 'DL4':
        file_path = os.path.join('/data2/zhengyu/workspace/flp', file_path)
    elif hostname == 'S4rawBer2y.local':
        file_path = os.path.join('/Users/zhengyu/Seafile/code/workspace/flp', file_path)
    elif hostname == 'DESKTOP-FTK3FVH':
        file_path = os.path.join('C:\\Users\\123\\Seafile\\code\\workspace\\flp', file_path)
    else:
        raise ValueError('Unknown hostname: {}'.format(socket.gethostname()))
    return file_path


class TupleSafeLoader(yaml.SafeLoader):
    def construct_python_tuple(self, node):
        return tuple(self.construct_sequence(node))


TupleSafeLoader.add_constructor(
    u'tag:yaml.org,2002:python/tuple',
    TupleSafeLoader.construct_python_tuple)


def load_yaml(file_path):
    file_path = get_file_path(file_path)
    files = glob.glob(file_path, recursive=True)
    assert(len(files) == 1)
    cfg = yaml.load(open(files[0], 'r'), Loader=TupleSafeLoader)
    return cfg


def load_pickle(file_path):
    file_path = get_file_path(file_path)
    files = glob.glob(file_path, recursive=True)
    assert(len(files) == 1)
    data = pickle.load(open(files[0], 'rb'))
    return data