4 files changed, 44 insertions, 12 deletions

diff --git a/context.c b/context.c
index 326c0c4..07d9304 100644
--- a/context.c
+++ b/context.c
@@ -66,12 +66,19 @@ static int Context_init(ContextObject *self, PyObject *args, PyObject *kwds)
 		return -1;
 	}
 
-	if (device && (rc = fdisk_assign_device(self->cxt, device, readonly)))
+	if (device && (rc = fdisk_assign_device(self->cxt, device, readonly))) {
+		set_PyErr_from_rc(-rc);
 		return -1;
-	if (details && (rc = fdisk_enable_details(self->cxt, details)))
+	}
+	if (details && (rc = fdisk_enable_details(self->cxt, details))) {
+		set_PyErr_from_rc(-rc);
 		return -1;
+	}
 
-	fdisk_get_partitions(self->cxt, &self->tb);
+	if ((fdisk_get_partitions(self->cxt, &self->tb) < 0)) {
+		set_PyErr_from_rc(-rc);
+		return -1;
+	}
 
 	return 0;
 }
@@ -82,6 +89,7 @@ static int Context_init(ContextObject *self, PyObject *args, PyObject *kwds)
 static PyObject *Context_assign_device(ContextObject *self, PyObject *args, PyObject *kwds)
 {
 	char *fname;
+	int rc;
 
 	if (!PyArg_ParseTuple(args, "s", &fname)) {
 		PyErr_SetString(PyExc_TypeError, ARG_ERR);
@@ -94,9 +102,11 @@ static PyObject *Context_assign_device(ContextObject *self, PyObject *args, PyOb
 	fdisk_assign_device(self->cxt, fname, 1);
 
 	self->lb = fdisk_get_label(self->cxt, NULL);
-	fdisk_get_partitions(self->cxt, &self->tb);
+	if ((rc = fdisk_get_partitions(self->cxt, &self->tb) < 0)) {
+		set_PyErr_from_rc(-rc);
+		return NULL;
+	}
 
-	/* XXX: check rc*/
 	Py_INCREF(Py_None);
 	return Py_None;
 }
diff --git a/fdisk.c b/fdisk.c
index d7d3fe4..d78e6e2 100644
--- a/fdisk.c
+++ b/fdisk.c
@@ -14,6 +14,25 @@
 
 #include "fdisk.h"
 
+void *set_PyErr_from_rc(int e)
+{
+	switch (e) {
+	case ENOMEM:
+		PyErr_SetString(PyExc_MemoryError, strerror(e));
+		break;
+	case EINVAL:
+		PyErr_SetString(PyExc_TypeError, strerror(e));
+		break;
+	case ENOSYS:
+		PyErr_SetString(PyExc_NotImplementedError, strerror(e));
+		break;
+	default:
+		PyErr_SetString(PyExc_Exception, strerror(e));
+	}
+
+	return NULL;
+}
+
 PyObject *PyObjectResultStr(const char *s)
 {
 	PyObject *result;
diff --git a/fdisk.h b/fdisk.h
index 39b4c95..3324aad 100644
--- a/fdisk.h
+++ b/fdisk.h
@@ -58,4 +58,6 @@ extern PyObject *PyObjectResultLabel(struct fdisk_label *lb);
 extern PyObject *PyObjectResultPartition(struct fdisk_partition *pa);
 extern PyObject *PyObjectResultPartType(struct fdisk_parttype *t);
 
+extern void *set_PyErr_from_rc(int err);
+
 #endif
diff --git a/partition.c b/partition.c
index 5284869..8685c3f 100644
--- a/partition.c
+++ b/partition.c
@@ -48,7 +48,8 @@ static int Partition_init(PartitionObject *self, PyObject *args, PyObject *kwds)
 	};
 	int partno_follow_default = 0,
 	    start_follow_default = 0,
-	    end_follow_default = 0;
+	    end_follow_default = 0,
+	    rc;
 
 	if (!PyArg_ParseTupleAndKeywords(args,
 					kwds, "|ppp", kwlist,
@@ -60,16 +61,16 @@ static int Partition_init(PartitionObject *self, PyObject *args, PyObject *kwds)
 	}
 
 	self->pa = fdisk_new_partition();
-	if (fdisk_partition_partno_follow_default(self->pa, partno_follow_default) < 0) {
-		PyErr_SetString(PyExc_RuntimeError, "Error setting partno_follow_default");
+	if ((rc = fdisk_partition_partno_follow_default(self->pa, partno_follow_default) < 0)) {
+		set_PyErr_from_rc(-rc);
 		return -1;
 	}
-	if (fdisk_partition_start_follow_default(self->pa, start_follow_default) < 0) {
-		PyErr_SetString(PyExc_RuntimeError, "Error setting start_follow_default");
+	if ((rc = fdisk_partition_start_follow_default(self->pa, start_follow_default) < 0)) {
+		set_PyErr_from_rc(-rc);
 		return -1;
 	}
-	if (fdisk_partition_end_follow_default(self->pa, end_follow_default) < 0) {
-		PyErr_SetString(PyExc_RuntimeError, "Error setting end_follow_default");
+	if ((rc = fdisk_partition_end_follow_default(self->pa, end_follow_default) < 0)) {
+		set_PyErr_from_rc(-rc);
 		return -1;
 	}