Fix reading database in newer version

[location/libloc.git] / src / python / writer.c

diff --git a/src/python/writer.c b/src/python/writer.c
index a15f77e8cffd8e9cb2ea92070bd329eb298a227f..00f7f6acfc59ab635eb6e48b9f5d684846c7ce46 100644 (file)
--- a/src/python/writer.c
+++ b/src/python/writer.c
@@ -21,6 +21,7 @@
 
 #include "locationmodule.h"
 #include "as.h"
+#include "country.h"
 #include "network.h"
 #include "writer.h"
 
@@ -38,12 +39,31 @@ static void Writer_dealloc(WriterObject* self) {
 }
 
 static int Writer_init(WriterObject* self, PyObject* args, PyObject* kwargs) {
-       // Create the writer object
-       int r = loc_writer_new(loc_ctx, &self->writer);
-       if (r)
+       PyObject* private_key = NULL;
+       FILE* f = NULL;
+
+       // Parse arguments
+       if (!PyArg_ParseTuple(args, "|O", &private_key))
                return -1;
 
-       return 0;
+       // Convert into FILE*
+       if (private_key && private_key != Py_None) {
+               int fd = PyObject_AsFileDescriptor(private_key);
+               if (fd < 0)
+                       return -1;
+
+               // Re-open file descriptor
+               f = fdopen(fd, "r");
+               if (!f) {
+                       PyErr_SetFromErrno(PyExc_IOError);
+                       return -1;
+               }
+       }
+
+       // Create the writer object
+       int r = loc_writer_new(loc_ctx, &self->writer, f);
+
+       return r;
 }
 
 static PyObject* Writer_get_vendor(WriterObject* self) {
@@ -118,6 +138,32 @@ static PyObject* Writer_add_as(WriterObject* self, PyObject* args) {
        return obj;
 }
 
+static PyObject* Writer_add_country(WriterObject* self, PyObject* args) {
+       struct loc_country* country;
+       const char* country_code;
+
+       if (!PyArg_ParseTuple(args, "s", &country_code))
+               return NULL;
+
+       // Create country object
+       int r = loc_writer_add_country(self->writer, &country, country_code);
+       if (r) {
+               switch (r) {
+                       case -EINVAL:
+                               PyErr_SetString(PyExc_ValueError, "Invalid network");
+                               break;
+
+                       default:
+                               return NULL;
+               }
+       }
+
+       PyObject* obj = new_country(&CountryType, country);
+       loc_country_unref(country);
+
+       return obj;
+}
+
 static PyObject* Writer_add_network(WriterObject* self, PyObject* args) {
        struct loc_network* network;
        const char* string = NULL;
@@ -149,17 +195,18 @@ static PyObject* Writer_add_network(WriterObject* self, PyObject* args) {
 
 static PyObject* Writer_write(WriterObject* self, PyObject* args) {
        const char* path = NULL;
+       int version = LOC_DATABASE_VERSION_UNSET;
 
-       if (!PyArg_ParseTuple(args, "s", &path))
+       if (!PyArg_ParseTuple(args, "s|i", &path, &version))
                return NULL;
 
-       FILE* f = fopen(path, "w");
+       FILE* f = fopen(path, "w+");
        if (!f) {
                PyErr_Format(PyExc_IOError, strerror(errno));
                return NULL;
        }
 
-       int r = loc_writer_write(self->writer, f);
+       int r = loc_writer_write(self->writer, f, (enum loc_database_version)version);
        fclose(f);
 
        // Raise any errors
@@ -178,6 +225,12 @@ static struct PyMethodDef Writer_methods[] = {
                METH_VARARGS,
                NULL,
        },
+       {
+               "add_country",
+               (PyCFunction)Writer_add_country,
+               METH_VARARGS,
+               NULL,
+       },
        {
                "add_network",
                (PyCFunction)Writer_add_network,
@@ -220,13 +273,13 @@ static struct PyGetSetDef Writer_getsetters[] = {
 
 PyTypeObject WriterType = {
        PyVarObject_HEAD_INIT(NULL, 0)
-       tp_name:                "location.Writer",
-       tp_basicsize:           sizeof(WriterObject),
-       tp_flags:               Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
-       tp_new:                 Writer_new,
-       tp_dealloc:             (destructor)Writer_dealloc,
-       tp_init:                (initproc)Writer_init,
-       tp_doc:                 "Writer object",
-       tp_methods:             Writer_methods,
-       tp_getset:              Writer_getsetters,
+       .tp_name =               "location.Writer",
+       .tp_basicsize =          sizeof(WriterObject),
+       .tp_flags =              Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
+       .tp_new =                Writer_new,
+       .tp_dealloc =            (destructor)Writer_dealloc,
+       .tp_init =               (initproc)Writer_init,
+       .tp_doc =                "Writer object",
+       .tp_methods =            Writer_methods,
+       .tp_getset =             Writer_getsetters,
 };