first viable definition for (optionally structured) Numpy array protobuf

Author: telamonian

.gitignore

@@ -0,0 +1,24 @@
+# test files stuff
+/figure/
+/out/
+
+# mac os stuff
+~$*
+.DS_Store
+
+# python stuff
+*.egg-info
+*.pyc
+*.ptx
+__pycache__/
+.ipynb_checkpoints/
+
+# jetbrains stuff
+.idea
+*.iml
+
+# cython builds
+cython**/*.c
+cython**/*.so
+cython**/build
+

npbuf/protobuf/dtype.proto

@@ -0,0 +1,51 @@
+syntax = "proto3";
+
+/**
+ * aiming for a minimal representation based on the numpy.dtype Python class.
+ * See https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html
+ */
+message Dtype {
+    enum ByteOrder {
+        LITTLE_ENDIAN = 0;    // '<' order
+        BIG_ENDIAN    = 1;    // '>' order
+        NATIVE        = 2;    // '=' order
+        NA            = 3;    // '|' order (not applicable)
+    }
+
+    // the labels of the Type enum are all parsable by np.dtype()
+    enum Type {
+        float64	   = 0;    // Double precision float: sign bit, 11 bits exponent, 52 bits mantissa
+        float32	   = 1;    // Single precision float: sign bit, 8 bits exponent, 23 bits mantissa
+        float16	   = 2;    // Half precision float: sign bit, 5 bits exponent, 10 bits mantissa
+        complex128 = 3;    // Complex number, represented by two 64-bit floats (real and imaginary components)
+        complex64  = 4;    // Complex number, represented by two 32-bit floats (real and imaginary components)
+        uint64	   = 5;    // Unsigned integer (0 to 18446744073709551615)
+        uint32	   = 6;    // Unsigned integer (0 to 4294967295)
+        uint16	   = 7;    // Unsigned integer (0 to 65535)
+        uint8	   = 8;    // Unsigned integer (0 to 255)
+        int64	   = 9;    // Integer (-9223372036854775808 to 9223372036854775807)
+        int32	   = 10;   // Integer (-2147483648 to 2147483647)
+        int16	   = 11;   // Integer (-32768 to 32767)
+        int8	   = 12;   // Byte (-128 to 127)
+        S128       = 13;   // String, 1 byte for every character, 128 bytes total
+        S64        = 14;   // String, 1 byte for every character, 64 bytes total
+        S32        = 15;   // String, 1 byte for every character, 32 bytes total
+        S16        = 16;   // String, 1 byte for every character, 16 bytes total
+        S8         = 17;   // String, 1 byte for every character, 8 bytes total
+    }
+
+    message SubDtype {
+        Type type = 1;
+        repeated int64 shape = 101;
+    }
+
+    Type type                     = 1;
+    ByteOrder byte_order          = 101;
+
+    // these fields are only used for structured dtypes
+    repeated string names = 201;        // Included for ordering of fields
+    map<string, Dtype> fields = 301;    // Order of map is not guaranteed
+
+    // only used if this dtype describes a sub-array
+    Dtype subdtype = 1001;
+}

npbuf/protobuf/ndarray.proto

@@ -0,0 +1,17 @@
+syntax = "proto3";
+
+import "dtype.proto";
+
+/**
+ * aiming for a minimal representation based the __array_interface__
+ * see https://docs.scipy.org/doc/numpy/reference/arrays.interface.html
+ */
+message NDArray {
+    repeated int64 shape = 1;
+    Dtype dtype = 101;
+    bytes data = 201;
+    repeated int64 strides = 301;
+
+    // TODO: enable compression
+    // optional bool compressed = X;
+}

setup.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+
+from Cython.Distutils import build_ext
+from distutils.errors import DistutilsExecError
+from glob import glob
+import os
+import setuptools
+from setuptools import find_packages, setup
+import shutil
+from subprocess import check_call
+
+# discover the path to this setup.py file
+thisScriptDir = os.path.dirname(os.path.realpath(__file__))
+
+class ProtocCommand(setuptools.Command):
+    user_options = [
+        ('protoc=', 'p', "(default: 'protoc') path to protoc compiler. The default value can be set via the PROTOC environment variable"),
+        ('raises=', 'r', '(default: True) if true, raise error on protobuf compilation error'),
+    ]
+
+    def initialize_options(self):
+        self.protoc = os.environ['PROTOC'] if 'PROTOC' in os.environ else 'protoc'
+        self.raises = True
+
+    def finalize_options(self):
+        print(f'searching for protoc executable at: {self.protoc}...')
+        self.protoc = shutil.which(self.protoc)
+        if self.protoc is None:
+            if self.raises:
+                raise DistutilsExecError('Could not find protoc executable.')
+            else:
+                print('Could not find protoc executable. Skipping protobuf compilation.')
+        else:
+            print(f'Found. Using protoc at: {self.protoc}')
+
+    def run(self):
+        if self.protoc is None:
+            # bail if self.raises is False and protoc isn't found
+            return
+
+        protoDir = os.path.join(thisScriptDir, 'npbuf', 'protobuf')
+        protoOutDir = os.path.join(thisScriptDir, 'npbuf', 'protobuf_py')
+        protoSrcs = glob(os.path.join(protoDir, '*.proto'))
+
+        # clean up any existing compiled protobufs
+        shutil.rmtree(protoOutDir, ignore_errors=True)
+        # create a new module for the compiled protobufs, including an `__init__.py`
+        os.mkdir(protoOutDir)
+        with open(os.path.join(protoOutDir, '__init__.py'), 'w') as f: pass
+
+        # compile protobuf files to .py python modules
+        protoc_python_cmd = [
+            self.protoc,
+            '--proto_path=%s' % protoDir,
+            '--python_out=%s' % protoOutDir,
+            ]
+        protoc_python_cmd.extend(protoSrcs)
+
+        check_call(protoc_python_cmd)
+
+class DevelopCommand(setuptools.command.develop.develop):
+    def run(self):
+        # pass options to the protoc command and run it
+        protocCommand = self.distribution.get_command_obj('protoc')
+        protocCommand.database = self.database
+        self.run_command('protoc')
+
+        # run the normal develop command
+        super().run(self)
+
+setup(
+    author = 'Max Klein',
+    cmdclass = {'build_ext': build_ext,
+                'develop': DevelopCommand,
+                'protoc': ProtocCommand
+    },
+    description = 'provides Python protobuf types that can be used to serialize/deserialize Numpy arrays',
+    license = 'Apache License, Version 2.0',
+    name = 'numpy-protobuf',
+    packages = find_packages(where='.', exclude=('npbuf.protobuf'))
+)