4 files changed, 27 insertions, 8 deletions
diff --git a/libfdisk/src/context.c b/libfdisk/src/context.c
index 4adc7cf30..908889b4f 100644
--- a/libfdisk/src/context.c
+++ b/libfdisk/src/context.c
@@ -49,6 +49,23 @@ struct fdisk_label *fdisk_context_get_label(struct fdisk_context *cxt, const cha
 	return NULL;
 }
 
+int __fdisk_context_switch_label(struct fdisk_context *cxt,
+				 struct fdisk_label *lb)
+{
+	if (!lb)
+		return -EINVAL;
+	cxt->label = lb;
+	DBG(LABEL, dbgprint("--> switching context to %s!", lb->name));
+	return 0;
+}
+
+int fdisk_context_switch_label(struct fdisk_context *cxt, const char *name)
+{
+	return __fdisk_context_switch_label(cxt,
+			fdisk_context_get_label(cxt, name));
+}
+
+
 static void reset_context(struct fdisk_context *cxt)
 {
 	size_t nlbs, i;
diff --git a/libfdisk/src/fdiskP.h b/libfdisk/src/fdiskP.h
index 00fc0cd54..104c67cb7 100644
--- a/libfdisk/src/fdiskP.h
+++ b/libfdisk/src/fdiskP.h
@@ -161,6 +161,7 @@ struct fdisk_label_operations {
 struct fdisk_label {
 	/* persistent information */
 	const char		*name;
+	enum fdisk_labeltype	id;		/* FDISK_DISKLABEL_* */
 	struct fdisk_parttype	*parttypes;
 	size_t			nparttypes;	/* number of items in parttypes[] */
 
@@ -190,8 +191,6 @@ struct fdisk_context {
 	unsigned long sector_size;	/* logical size */
 	unsigned long alignment_offset;
 
-	enum fdisk_labeltype disklabel;	/* current disklabel */
-
 	/* alignment */
 	unsigned long grain;		/* alignment unit */
 	sector_t first_lba;		/* recommended begin of the first partition */
@@ -207,7 +206,9 @@ struct fdisk_context {
 					 * FIXME: use any enum rather than hardcoded number */
 };
 
-
+/* context.c */
+extern int __fdisk_context_switch_label(struct fdisk_context *cxt,
+				    struct fdisk_label *lb);
 
 /* alignment.c */
 extern sector_t fdisk_topology_get_first_lba(struct fdisk_context *cxt);
diff --git a/libfdisk/src/label.c b/libfdisk/src/label.c
index 57b756591..e0ec9022f 100644
--- a/libfdisk/src/label.c
+++ b/libfdisk/src/label.c
@@ -8,7 +8,6 @@ int fdisk_probe_labels(struct fdisk_context *cxt)
 {
 	size_t i;
 
-	cxt->disklabel = FDISK_DISKLABEL_ANY;
 	cxt->label = NULL;
 
 	for (i = 0; i < cxt->nlabels; i++) {
@@ -24,8 +23,7 @@ int fdisk_probe_labels(struct fdisk_context *cxt)
 			continue;
 		}
 
-		cxt->label = lb;
-		DBG(LABEL, dbgprint("detected a %s label", lb->name));
+		__fdisk_context_switch_label(cxt, lb);
 		return 0;
 	}
 
@@ -41,7 +39,7 @@ int fdisk_probe_labels(struct fdisk_context *cxt)
  */
 int fdisk_dev_has_disklabel(struct fdisk_context *cxt)
 {
-	return cxt && cxt->disklabel != FDISK_DISKLABEL_ANY;
+	return cxt && cxt->label;
 }
 
 /**
@@ -53,7 +51,7 @@ int fdisk_dev_has_disklabel(struct fdisk_context *cxt)
  */
 int fdisk_dev_is_disklabel(struct fdisk_context *cxt, enum fdisk_labeltype l)
 {
-	return cxt && cxt->disklabel == l;
+	return cxt && cxt->label && cxt->label->id == l;
 }
 
 /**
diff --git a/libfdisk/src/libfdisk.h b/libfdisk/src/libfdisk.h
index 927587622..7fcd3d74e 100644
--- a/libfdisk/src/libfdisk.h
+++ b/libfdisk/src/libfdisk.h
@@ -56,6 +56,9 @@ extern int fdisk_context_assign_device(struct fdisk_context *cxt,
 extern struct fdisk_label *fdisk_context_get_label(struct fdisk_context *cxt,
 				const char *name);
 
+extern int fdisk_context_switch_label(struct fdisk_context *cxt,
+				const char *name);
+
 /* parttype.c */
 extern struct fdisk_parttype *fdisk_get_parttype_from_code(struct fdisk_context *cxt,
                                 unsigned int code);