Revision 331aa0ec image_creator/disk.py

b/image_creator/disk.py
33 33

  
34 34
from image_creator.util import get_command
35 35
from image_creator.util import warn, progress, success, output, FatalError
36

  
36
from image_creator.gpt import GPTPartitionTable
37 37
import stat
38 38
import os
39 39
import tempfile
......
200 200
            raise FatalError("Multiple operating systems found."
201 201
                            "We only support images with one filesystem.")
202 202
        self.root = roots[0]
203
        self.gdev = self.g.part_to_dev(self.root)
204
        self.parttype = self.g.part_get_parttype(self.gdev)
205

  
203 206
        self.ostype = self.g.inspect_get_type(self.root)
204 207
        self.distro = self.g.inspect_get_distro(self.root)
205 208
        success('found a(n) %s system' % self.distro)
......
256 259
        """
257 260
        output("Shrinking image (this may take a while)...", False)
258 261

  
259
        dev = self.g.part_to_dev(self.root)
260
        parttype = self.g.part_get_parttype(dev)
261
        if parttype != 'msdos':
262
        if self.parttype not in 'msdos' 'gpt':
262 263
            raise FatalError("You have a %s partition table. "
263
                "Only msdos partitions are supported" % parttype)
264
                "Only msdos and gpt partitions are supported" % self.parttype)
264 265

  
265
        last_partition = self.g.part_list(dev)[-1]
266
        last_partition = self.g.part_list(self.gdev)[-1]
266 267

  
267 268
        if last_partition['part_num'] > 4:
268 269
            raise FatalError("This disk contains logical partitions. "
269 270
                "Only primary partitions are supported.")
270 271

  
271
        part_dev = "%s%d" % (dev, last_partition['part_num'])
272
        part_dev = "%s%d" % (self.gdev, last_partition['part_num'])
272 273
        fs_type = self.g.vfs_type(part_dev)
273 274
        if not re.match("ext[234]", fs_type):
274 275
            warn("Don't know how to resize %s partitions." % vfs_type)
275
            return
276
            return self.size()
276 277

  
277 278
        self.g.e2fsck_f(part_dev)
278 279
        self.g.resize2fs_M(part_dev)
......
283 284
        block_cnt = int(
284 285
            filter(lambda x: x[0] == 'Block count', out)[0][1])
285 286

  
286
        sector_size = self.g.blockdev_getss(dev)
287
        sector_size = self.g.blockdev_getss(self.gdev)
287 288

  
288 289
        start = last_partition['part_start'] / sector_size
289 290
        end = start + (block_size * block_cnt) / sector_size - 1
290 291

  
291
        self.g.part_del(dev, last_partition['part_num'])
292
        self.g.part_add(dev, 'p', start, end)
292
        self.g.part_del(self.gdev, last_partition['part_num'])
293
        self.g.part_add(self.gdev, 'p', start, end)
293 294

  
294 295
        new_size = (end + 1) * sector_size
295 296
        success("new image size is %dMB" %
296 297
                            ((new_size + 2 ** 20 - 1) // 2 ** 20))
298

  
299
        if self.parttype == 'gpt':
300
            ptable = GPTPartitionTable(self.device)
301
            return ptable.shrink(new_size)
302

  
297 303
        return new_size
298 304

  
299 305
    def size(self):
......
302 308
        The size returned by this method is the size of the space occupied by
303 309
        the partitions (including the space before the first partition).
304 310
        """
305
        dev = self.g.part_to_dev(self.root)
306
        last = self.g.part_list(dev)[-1]
307 311

  
308
        return last['part_end'] + 1
312
        if self.parttype == 'msdos':
313
            dev = self.g.part_to_dev(self.root)
314
            last = self.g.part_list(dev)[-1]
315
            return last['part_end'] + 1
316
        elif self.parttype == 'gpt':
317
            ptable = GPTPartitionTable(self.device)
318
            return ptable.size()
319
        else:
320
            raise FatalError("Unsupported partition table type: %s" % parttype)
309 321

  
310 322
    def dump(self, outfile):
311 323
        """Dumps the content of device into a file.

Also available in: Unified diff