Довольно отформатированный кодировщик numpy JSON

Я искал способ сохранить данные numpy с помощью json, сохранив удобочитаемый формат печати numpy.

Вдохновленный этим ответом, я решил использовать pprint вместо base64 для записи данных с желаемым форматированием, так что дано:

import numpy as np
data = np.random.random((1,3,2))

Полученный файл на диске должен выглядеть так:

{
    "__dtype__": "float64", 
    "__ndarray__": [[[0.7672818918130646 , 0.6846412220229668 ],
                     [0.7082023466738064 , 0.0896531267221291 ],
                     [0.43898454934160147, 0.9245898883694668 ]]]
}

Появилось несколько заиканий.

• Хотя json мог считывать списки списков, отформатированных как [[...]], у него были проблемы с форматированием с плавающей запятой numpy. Например, [[0., 0., 0.]] вызовет ошибку при обратном считывании, а [[0.0, 0.0, 0.0]] подойдет.

• pformat выведет array([[0., 0., 0.]]), где array() должно быть проанализировано, иначе json выдаст ошибку при обратном считывании данных.

Чтобы исправить это, мне пришлось выполнить синтаксический анализ строк, который привел к моему текущему коду ниже:

import json, sys
import numpy as np
import pprint as pp

# Set numpy's printoptions to display all the data with max precision
np.set_printoptions(threshold=np.inf,
                    linewidth=sys.maxsize,
                    suppress=True,
                    nanstr='0.0',
                    infstr='0.0', 
                    precision=np.finfo(np.longdouble).precision)     



# Modified version of Adam Hughes's https://stackoverflow.com/a/27948073/1429402
def save_formatted(fname,data):

    class NumpyEncoder(json.JSONEncoder):
        def default(self, obj):
            if isinstance(obj, np.ndarray):
                return {'__ndarray__': self.numpy_to_string(obj),
                        '__dtype__': str(obj.dtype)}            

            return json.JSONEncoder.default(self, obj)


        def numpy_to_string(self,data):
            ''' Use pprint to generate a nicely formatted string
            '''

            # Get rid of array(...) and keep only [[...]]
            f = pp.pformat(data, width=sys.maxsize)
            f = f[6:-1].splitlines() # get rid of array(...) and keep only [[...]]

            # Remove identation caused by printing "array(" 
            for i in xrange(1,len(f)):
                f[i] = f[i][6:]

            return '\n'.join(f)


    # Parse json stream and fix formatting.
    # JSON doesn't support float arrays written as [0., 0., 0.]
    # so we look for the problematic numpy print syntax and correct
    # it to be readable natively by JSON, in this case: [0.0, 0.0, 0.0]
    with open(fname,'w') as io:
        for line in json.dumps(data, sort_keys=False, indent=4, cls=NumpyEncoder).splitlines():
            if '"__ndarray__": "' in line:
                index = line.index('"__ndarray__": "')
                lines = line.split('"__ndarray__": "')[-1][:-1]
                lines = lines.replace('. ','.0')  # convert occurences of ". " to ".0"    ex: 3. , 2. ]
                lines = lines.replace('.,','.0,') # convert occurences of ".," to ".0,"   ex: 3., 2.,
                lines = lines.replace('.]','.0]') # convert occurences of ".]" to ".0],"  ex: 3., 2.]
                lines = lines.split('\\n')

                # write each lines with appropriate indentation
                for i in xrange(len(lines)):
                    if i == 0:
                        indent = ' '*index
                        io.write(('%s"__ndarray__": %s\n"'%(indent,lines[i]))[:-1]) 
                    else:
                        indent = ' '*(index+len('"__ndarray__": "')-1)
                        io.write('%s%s\n'%(indent,lines[i]))                        

            else:
                io.write('%s\n'%line)



def load_formatted(fname):

    def json_numpy_obj_hook(dct):
        if isinstance(dct, dict) and '__ndarray__' in dct:
            return np.array(dct['__ndarray__']).astype(dct['__dtype__'])        
        return dct

    with open(fname,'r') as io:
        return json.load(io, object_hook=json_numpy_obj_hook)

Тестировать:

data = np.random.random((200,3,1000))
save_formatted('test.data', data)
data_ = load_formatted('test.data')

print np.allclose(data,data_) # Returns True

ВОПРОС

Мое решение меня устраивает, но его аспект синтаксического анализа строк замедляет работу с большими массивами данных. Есть ли лучший способ добиться желаемого эффекта? Может ли regular expression заменить мою последовательность вызовов str.replace()? Или, может быть, pprint можно использовать для правильного форматирования моей строки в первую очередь? Есть ли лучший способ сделать списки записи json такими же, как форматирование печати numpy?