3
\$\begingroup\$

I'm learning python and for educational purposes I implemented tar archive parser. I'm not beginner programmer. I would like to receive some feedback and tips about code, what can I improve, what could be done better and so on.

Implementation:

tar.py

#!/usr/bin/env python3

import io
import os
import sys
import math
import json

class Tar:

    BLOCK_SIZE = 512

    def __init__(self, file_path):
        if not file_path or len(file_path) == 0:
            raise ValueError("Bad file path")
        self.file_path = file_path

    def __enter__(self):
        self.input_stream = open(self.file_path, "rb")
        self.headers = []
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def close(self):
        if self.input_stream:
            self.input_stream.close()

    def get_all_files(self):
        self.__scan()
        return list(map(
            lambda f: FileSnapshot(f.file_name, f.file_size, f.file_mode, f.flag), 
            self.headers
        ))

    def extract_file(self, file_name, target_folder=os.getcwd()):
        if not file_name or len(file_name) == 0:
            raise ValueError("Bad file name")
        if not target_folder or len(target_folder) == 0:
            raise ValueError("Bad target folder")
        self.__scan()
        result = list(filter(
            lambda fh: fh.flag == 0 and fh.file_name == file_name, 
            self.headers
        ))
        if len(result) == 0:
            raise RuntimeError("File '{}' not found".format(file_name))
        fh = result[0]
        leaf = os.path.basename(fh.file_name)
        f_path = os.path.join(target_folder, leaf)
        self.__extract(fh, f_path)

    def extract_all(self, target_folder=os.getcwd()):
        if not target_folder or len(target_folder) == 0:
            raise ValueError("Bad target folder")
        self.__scan()
        for fh in self.headers:
            f_path = os.path.join(target_folder, fh.file_name)
            if fh.flag == 5: # if directory
                os.makedirs(f_path, exist_ok=True)
            elif fh.flag == 0: # if regular file
                parent = os.path.dirname(os.path.abspath(f_path))
                os.makedirs(parent, exist_ok=True)
                self.__extract(fh, f_path)


    def __extract(self, fh, file_name):
        with open(file_name, "wb") as f:
            if fh.file_size > 0:
                total = 0
                bytes_left = fh.file_size
                self.input_stream.seek(fh.offset, 0)
                while bytes_left > 0:
                    data = self.input_stream.read(Tar.BLOCK_SIZE)
                    data = data[:bytes_left]
                    f.write(data)
                    bytes_left -= len(data)

    def __scan(self): # iterate over headers
        if len(self.headers) == 0:
            while True:
                block = self.input_stream.read(Tar.BLOCK_SIZE)
                if len(block) < Tar.BLOCK_SIZE:
                    break
                h = self.__get_file_header(block)
                if not len(h.magic) > 0:
                    break
                # ommit regular file bytes
                if h.flag == 0:
                    h.set_offset(self.input_stream.tell())
                    if h.file_size > 0:
                        if h.file_size % Tar.BLOCK_SIZE != 0:
                            bytes_to_skeep = math.ceil(h.file_size / Tar.BLOCK_SIZE) * Tar.BLOCK_SIZE
                        else:
                            bytes_to_skeep = h.file_size
                        self.input_stream.seek(bytes_to_skeep, 1)
                self.headers.append(h)


    def __get_file_header(self, block):
        try:
            file_name = self.__get_file_name(block)
            file_mode = self.__get_file_mode(block)
            uid = self.__get_uid(block)
            gid = self.__get_gid(block)
            file_size = self.__get_file_size(block)
            mtime = self.__get_mtime(block)
            chksum = self.__get_chksum(block)
            type_flag = self.__get_type_flag(block)
            linkname = self.__get_linkname(block)
            magic = self.__get_magic(block)
            version = self.__get_version(block)
            uname = self.__get_uname(block)
            gname = self.__get_gname(block)
            devmajor = self.__get_devmajor(block)
            devminor = self.__get_devminor(block)
            prefix = self.__get_prefix(block)
        except Exception as e:
            raise RuntimeError("Broken file") from e
        header = FileHeader(file_name, file_size, file_mode, uid, gid, 
            mtime, chksum, type_flag, linkname, magic, version, 
            uname, gname, devmajor, devminor, prefix)
        return header


    def __get_file_name(self, block): # string
        offset, size = 0, 100
        fname = self.__get_block_data(block, offset, size)
        fname = fname[0:fname.find(b'\x00')].decode().strip()
        return fname

    def __get_file_mode(self, block): # string
        offset, size = 100, 8
        mode = self.__get_block_data(block, offset, size)
        mode = mode[:mode.find(b'\x00')].decode().strip()
        return mode

    def __get_uid(self, block): # string
        offset, size = 108, 8
        uid = self.__get_block_data(block, offset, size)
        uid = uid[:uid.find(b'\x00')].decode().strip()
        return uid

    def __get_gid(self, block): # string
        offset, size = 116, 8
        gid = self.__get_block_data(block, offset, size)
        gid = gid[:gid.find(b'\x00')].decode().strip()
        return gid

    def __get_file_size(self, block): # int
        offset, size = 124, 12
        size = self.__get_block_data(block, offset, size)
        size = size[:size.find(b'\x00')].decode().strip()
        if len(size) > 0:
            size = int(size, 8)
        else:
            size = 0
        return size

    def __get_mtime(self, block): # int
        offset, size = 136, 12
        mtime = self.__get_block_data(block, offset, size)
        mtime = mtime[:len(mtime)-1]
        mtime = mtime[:mtime.find(b'\x00')].decode().strip()
        if len(mtime) > 0:
            mtime = int(mtime, 8)
        else:
            mtime = 0
        return mtime

    def __get_chksum(self, block): # int
        offset, size = 148, 8
        chksum = self.__get_block_data(block, offset, size)
        chksum = chksum[:chksum.find(b'\x00')].decode().strip()
        if len(chksum) > 0:
            chksum = int(chksum)
        else:
            chksum = 0
        return chksum

    def __get_type_flag(self, block): # int
        offset, size = 156, 1
        flag = self.__get_block_data(block, offset, size)
        if flag == b'\x00':
            flag = 0
        elif flag == b'x':
            flag = 11
        else:
            flag = int(flag)
        return flag

    def __get_linkname(self, block): # string (applicable if type_flag = 1 or 2)
        offset, size = 157, 100
        linkname = self.__get_block_data(block, offset, size)
        return linkname[:linkname.find(b'\x00')].decode().strip()

    def __get_magic(self, block): # string
        offset, size = 257, 6
        magic = self.__get_block_data(block, offset, size)
        magic = magic[:magic.find(b'\x00')].decode().strip()
        return magic

    def __get_version(self, block): # string
        offset, size = 263, 2
        version = self.__get_block_data(block, offset, size)
        version = version[:len(version)-1].decode().strip()
        return version

    def __get_uname(self, block): # string
        offset, size = 265, 32
        uname = self.__get_block_data(block, offset, size)
        uname = uname[:uname.find(b'\x00')].decode().strip()
        return uname

    def __get_gname(self, block): # string
        offset, size = 297, 32
        gname = self.__get_block_data(block, offset, size)
        gname = gname[:gname.find(b'\x00')].decode().strip()
        return gname

    def __get_devmajor(self, block): # string
        offset, size = 329, 8
        devmajor = self.__get_block_data(block, offset, size)
        devmajor = devmajor[:devmajor.find(b'\x00')].decode().strip()
        return devmajor

    def __get_devminor(self, block): # string
        offset, size = 337, 8
        devminor = self.__get_block_data(block, offset, size)
        devminor = devminor[:devminor.find(b'\x00')].decode().strip()
        return devminor

    def __get_prefix(self, block): # string
        offset, size = 345, 155
        prefix = self.__get_block_data(block, offset, size)
        prefix = prefix[:prefix.find(b'\x00')].decode().strip()
        return prefix

    def __get_block_data(self, block, offset, size):
        return block[offset:offset+size]

class FileSnapshot:
    def __init__(self, file_name, file_size, file_mode, flag):
        self.file_name = file_name
        self.file_size = file_size
        self.file_mode = file_mode
        self.flag = flag

    def __repr__(self):
        return self.file_name

class FileHeader:
    def __init__(self, file_name, file_size, file_mode, uid, gid, mtime, 
    chksum, flag, linkname, magic, version, uname, gname, devmajor, devminor, prefix):
        self.file_name = file_name
        self.file_size = file_size
        self.file_mode = file_mode
        self.uid = uid
        self.gid = gid
        self.mtime = mtime
        self.chksum = chksum
        self.flag = flag
        self.linkname = linkname
        self.magic = magic
        self.version = version
        self.uname = uname
        self.gname = gname
        self.devmajor = devmajor
        self.devminor = devminor
        self.prefix = prefix

    def set_offset(self, offset):
        self.offset = offset 


def usage():
    u = """
    Usage:
    tar.py <archive.tar> --list                 List all files in the archive
    tar.py <archive.tar> --extract-all          Extract all files from the archive
    tar.py <archive.tar> --extract <file>       Extract single file from the archive
    """
    print(u)
    sys.exit(1)

if __name__ == "__main__":
    try:
        if len(sys.argv) > 2:
            archive = sys.argv[1]
            operation = sys.argv[2]
            with Tar(archive) as t:
                if operation == "--list":
                    files = t.get_all_files()
                    for f in files:
                        print(f)
                elif operation == "--extract-all":
                    t.extract_all()
                elif operation == "--extract":
                    if len(sys.argv) > 3:
                        file_name = sys.argv[3]
                        t.extract_file(file_name)
                    else:
                        usage()
        else:
            usage()
    except Exception as e:
        print("Error: {}".format(str(e)))
        sys.exit(1)

tartest.py

#!/usr/bin/env python3

import unittest
import tar
import os

class TarTest(unittest.TestCase):

    def test_get_all_files(self):
        # given
        with tar.Tar("tartest.tar") as t:

            # when
            files = t.get_all_files()

            # then
            self.assertTrue(len(files) == 5)
            self.assertTrue(self.containsFile(files, "tartest/a.txt"))
            self.assertTrue(self.containsFile(files, "tartest/b.txt"))
            self.assertTrue(self.containsFile(files, "tartest/foo/c.txt"))

    def test_extract_file(self):
        # given
        with tar.Tar("tartest.tar") as t:

            # when
            t.extract_file("tartest/a.txt")
            t.extract_file("tartest/foo/c.txt")

            # then
            self.assertTrue(os.path.isfile("a.txt"))
            self.assertTrue(self.fileContains("a.txt", "This is file a"))

            self.assertTrue(os.path.isfile("c.txt"))
            self.assertTrue(self.fileContains("c.txt", "This is file c"))

            os.remove("a.txt")
            os.remove("c.txt")

    def test_extract_all(self):
        # given
        with tar.Tar("tartest.tar") as t:

            # when
            t.extract_all()

            # then
            self.assertTrue(os.path.isdir("tartest"))
            self.assertTrue(os.path.isdir("tartest/foo"))
            self.assertTrue(os.path.isfile("tartest/a.txt"))
            self.assertTrue(os.path.isfile("tartest/b.txt"))
            self.assertTrue(os.path.isfile("tartest/foo/c.txt"))

            os.system("rm -rf tartest")

    def containsFile(self, files, file_name):
        for f in files:
            if f.file_name == file_name:
                return True
        return False

    def fileContains(self, file_name, content):
        with open(file_name) as f:
            return content == f.read().splitlines()[0]

if __name__ == '__main__':
    unittest.main()
\$\endgroup\$
  • \$\begingroup\$ Oh, one last thing: I tried using your code to unzip a file tar'ed with 7zip, and it failed trying to convert the flag (b'L', which is a vendor flag) to an int; you might want your code to fail more gracefully in the case of an unrecognized flag (and use a str flag instead of an int one) \$\endgroup\$ – LimeHunter7 Mar 26 at 19:44
4
\$\begingroup\$

I can offer some tips for making your code more pythonic.

General

  • It's good practice to add a docstring for each module, class, and documented function.
  • Imports io and json are unused.
  • In the Tar.__extract method, the variable total is unused.

class Tar

  • not file_path or len(file_path) == 0: If the user inputs an empty string, not file_path is sufficient (and None is not a possible value unless input manually). More to the point, you are not exactly detecting a "bad file path". You could use os.path.exists for a more robust check. Alternatively, don't validate the path at all, and consider a try... except OSError block in your __enter__ method; this will avoid race conditions. (You perform similar checks in extract_file and extract_all that can also be changed.)

  • You have an __enter__ method and an __exit__ method, which allows your class to be used with a context-manager, excellent! However, you also provide a close function, without providing a corresponding open function, which means that close could never be reasonably called by the user. Eliminate close or add open.

  • You invoke name-mangling by using double-underscores on methods like __extract; this is fine to prevent truly "private" data members from clashing with those from a superclass or subclass, but on methods it makes inheriting from your class (say, to extend it with logging features) unnecessarily difficult. To mark a member as "private", a single leading underscore is enough.

  • Similarly, in the interest of being able to subclass your class, you should consider self.BLOCK_SIZE instead of Tar.BLOCK_SIZE (though maybe this is a constant of the tar format?).

  • list(map(...)): It's generally more clear to replace this with a list comprehension (and as opposed to lambda, sometimes more performant):

    def get_all_files(self):
        self._scan()
        return [FileSnapshot(f.file_name, f.file_size, f.file_mode, f.flag) for f in self.headers]
    
  • list(filter(...)): To get the first match, it's generally better to use a generator comprehension:

    def extract_file(...):
        ...
        try:
            result = next(fh for fh in self.headers if fh.flag == 0 and fh.file_name == file_name)
        except StopIteration:
            raise RuntimeError("File '{}' not found".format(file_name))
        ...
    

class FileSnapshot, class FileHeader

  • There is a lot of boilerplate code here, that could be eliminated with e.g. the @dataclass decorator.

    from dataclasses import dataclass
    ...
    @dataclass
    class FileSnapshot:
        file_name : str
        file_size : int
        ...
    
  • __repr__ methods are generally supposed to return code that will reproduce the object; consider renaming this method to __str__ instead.

__main__

  • Take advantage of the standard library argparse module. For example, it makes extending your --extract switch to extract multiple files easier, provides error checking and usage strings, and can be used to initialize archive as a Tar automatically.

    from argparse import ArgumentParser
    ...
    if __name__ == '__main__':
        parser = ArgumentParser(description='.tar archive extractor')
        parser.add_argument('archive', type=Tar, help='...')
        group = parser.add_mutually_exclusive_group(required=True)
        group.add_argument('--list', action='store_true', help='List files')
        group.add_argument('--extract-all', action='store_true', help='Extract all')
        group.add_argument('--extract', nargs='+', help='Extract some')
        args = parser.parse_args()
        with args.archive as t:
            ...
    

Code

Here's my take on your code

#!/usr/bin/env python3
'''TODO: docstring'''

import os
import math
from dataclasses import dataclass

class Tar:
    '''TODO docstring'''
    BLOCK_SIZE = 512

    def __init__(self, file_path):
        self.file_path = file_path
        self.input_stream = None
        self.headers = []

    def __enter__(self):
        self.input_stream = open(self.file_path, "rb")
        self.headers = []
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.input_stream is not None:
            self.input_stream.close()

    def get_all_files(self):
        '''TODO docstring'''
        self._scan()
        return [FileSnapshot(f.file_name, f.file_size, f.file_mode, f.flag) 
                for f in self.headers]

    def extract_file(self, file_name, target_folder=os.getcwd()):
        '''TODO docstring'''
        self._scan()
        try:
            fh = next(fh for fh in self.headers if fh.flag == 0 and fh.file_name == file_name)
        except StopIteration:
            raise RuntimeError("File '{}' not found".format(file_name))
        leaf = os.path.basename(fh.file_name)
        f_path = os.path.join(target_folder, leaf)
        self._extract(fh, f_path)

    def extract_all(self, target_folder=os.getcwd()):
        '''TODO docstring'''
        self._scan()
        for fh in self.headers:
            f_path = os.path.join(target_folder, fh.file_name)
            if fh.flag == 5: # if directory
                os.makedirs(f_path, exist_ok=True)
            elif fh.flag == 0: # if regular file
                parent = os.path.dirname(os.path.abspath(f_path))
                os.makedirs(parent, exist_ok=True)
                self._extract(fh, f_path)


    def _extract(self, fh, file_name):
        with open(file_name, "wb") as f:
            if fh.file_size > 0:
                bytes_left = fh.file_size
                self.input_stream.seek(fh.offset, 0)
                while bytes_left > 0:
                    data = self.input_stream.read(Tar.BLOCK_SIZE)
                    data = data[:bytes_left]
                    f.write(data)
                    bytes_left -= len(data)

    def _scan(self): # iterate over headers
        if len(self.headers) == 0:
            while True:
                block = self.input_stream.read(Tar.BLOCK_SIZE)
                if len(block) < Tar.BLOCK_SIZE:
                    break
                h = self._get_file_header(block)
                if not len(h.magic) > 0:
                    break
                # omit regular file bytes
                if h.flag == 0:
                    h.offset = self.input_stream.tell()
                    if h.file_size > 0:
                        if h.file_size % Tar.BLOCK_SIZE != 0:
                            bytes_to_skeep = math.ceil(h.file_size / Tar.BLOCK_SIZE) * Tar.BLOCK_SIZE
                        else:
                            bytes_to_skeep = h.file_size
                        self.input_stream.seek(bytes_to_skeep, 1)
                self.headers.append(h)


    def _get_file_header(self, block):
        try:
            return FileHeader(
                self._get_file_name(block),
                self._get_file_size(block),
                self._get_file_mode(block),
                self._get_uid(block),
                self._get_gid(block),
                self._get_mtime(block),
                self._get_chksum(block),
                self._get_type_flag(block),
                self._get_linkname(block),
                self._get_magic(block),
                self._get_version(block),
                self._get_uname(block),
                self._get_gname(block),
                self._get_devmajor(block),
                self._get_devminor(block),
                self._get_prefix(block)
            )
        except Exception as e:
            raise RuntimeError("Broken file") from e


    def _get_file_name(self, block): # string
        offset, size = 0, 100
        fname = self._get_block_data(block, offset, size)
        fname = fname[0:fname.find(b'\x00')].decode().strip()
        return fname

    def _get_file_mode(self, block): # string
        offset, size = 100, 8
        mode = self._get_block_data(block, offset, size)
        mode = mode[:mode.find(b'\x00')].decode().strip()
        return mode

    def _get_uid(self, block): # string
        offset, size = 108, 8
        uid = self._get_block_data(block, offset, size)
        uid = uid[:uid.find(b'\x00')].decode().strip()
        return uid

    def _get_gid(self, block): # string
        offset, size = 116, 8
        gid = self._get_block_data(block, offset, size)
        gid = gid[:gid.find(b'\x00')].decode().strip()
        return gid

    def _get_file_size(self, block): # int
        offset, size = 124, 12
        size = self._get_block_data(block, offset, size)
        size = size[:size.find(b'\x00')].decode().strip()
        if len(size) > 0:
            size = int(size, 8)
        else:
            size = 0
        return size

    def _get_mtime(self, block): # int
        offset, size = 136, 12
        mtime = self._get_block_data(block, offset, size)
        mtime = mtime[:len(mtime)-1]
        mtime = mtime[:mtime.find(b'\x00')].decode().strip()
        if len(mtime) > 0:
            mtime = int(mtime, 8)
        else:
            mtime = 0
        return mtime

    def _get_chksum(self, block): # int
        offset, size = 148, 8
        chksum = self._get_block_data(block, offset, size)
        chksum = chksum[:chksum.find(b'\x00')].decode().strip()
        if len(chksum) > 0:
            chksum = int(chksum)
        else:
            chksum = 0
        return chksum

    def _get_type_flag(self, block): # int
        offset, size = 156, 1
        flag = self._get_block_data(block, offset, size)
        if flag == b'\x00':
            flag = 0
        elif flag == b'x':
            flag = 11
        else:
            flag = int(flag)
        return flag

    def _get_linkname(self, block): # string (applicable if type_flag = 1 or 2)
        offset, size = 157, 100
        linkname = self._get_block_data(block, offset, size)
        return linkname[:linkname.find(b'\x00')].decode().strip()

    def _get_magic(self, block): # string
        offset, size = 257, 6
        magic = self._get_block_data(block, offset, size)
        magic = magic[:magic.find(b'\x00')].decode().strip()
        return magic

    def _get_version(self, block): # string
        offset, size = 263, 2
        version = self._get_block_data(block, offset, size)
        version = version[:len(version)-1].decode().strip()
        return version

    def _get_uname(self, block): # string
        offset, size = 265, 32
        uname = self._get_block_data(block, offset, size)
        uname = uname[:uname.find(b'\x00')].decode().strip()
        return uname

    def _get_gname(self, block): # string
        offset, size = 297, 32
        gname = self._get_block_data(block, offset, size)
        gname = gname[:gname.find(b'\x00')].decode().strip()
        return gname

    def _get_devmajor(self, block): # string
        offset, size = 329, 8
        devmajor = self._get_block_data(block, offset, size)
        devmajor = devmajor[:devmajor.find(b'\x00')].decode().strip()
        return devmajor

    def _get_devminor(self, block): # string
        offset, size = 337, 8
        devminor = self._get_block_data(block, offset, size)
        devminor = devminor[:devminor.find(b'\x00')].decode().strip()
        return devminor

    def _get_prefix(self, block): # string
        offset, size = 345, 155
        prefix = self._get_block_data(block, offset, size)
        prefix = prefix[:prefix.find(b'\x00')].decode().strip()
        return prefix

    def _get_block_data(self, block, offset, size):
        return block[offset:offset+size]

@dataclass
class FileSnapshot:
    '''TODO: docstring'''
    file_name: str
    file_size: int
    file_mode: str
    flag: int

    def __str__(self):
        return self.file_name

@dataclass
class FileHeader:
    '''TODO: docstring'''
    file_name: str
    file_size: int
    file_mode: str
    uid: str
    gid: str
    mtime: int
    chksum: int
    flag: int
    linkname: str 
    magic: str
    version: str
    uname: str
    gname: str
    devmajor: str
    devminor: str
    prefix: str
    offset: int = 0

if __name__ == "__main__":
    def main():
        from argparse import ArgumentParser
        parser = ArgumentParser(description='.tar archive extractor')
        parser.add_argument('archive', type=Tar,
                            help='The tar archive file')
        group = parser.add_mutually_exclusive_group(required=True)
        group.add_argument('--list', action='store_true', 
                           help='List all files in the archive')
        group.add_argument('--extract-all', action='store_true', 
                           help='Extract all files from the archive')
        group.add_argument('--extract', nargs='+', dest='files',
                           help='Extract specified files from the archive')
        args = parser.parse_args()
        with args.archive as t:
            if args.list:
                files = t.get_all_files()
                for file in files:
                    print(file)
            elif args.extract_all:
                t.extract_all()
            else:
                for file in args.files:
                    t.extract_file(file)

    main()

| improve this answer | |
\$\endgroup\$
  • 1
    \$\begingroup\$ All the _get_* could be probably done easier by having a _get_block_data that truncates to the null byte and decodes and strips the value and then you can have e.g. _get_linkname = functools.partialmethod(_get_block_data, offset=157, size=100). This way you only need to spell out the two where you need to do additional checking. Or maybe even better, save the sizes in a class variable and iterate over the block in chunks of varying sizes. \$\endgroup\$ – Graipher Mar 27 at 11:27
  • 1
    \$\begingroup\$ It's true, the get methods are not very DRY; even just adding the find/slice/decode/strip portion to _get_block_data would save on boilerplate, or maybe factoring it into a _get_block_str since _get_type_flag exceptionally doesn't use it \$\endgroup\$ – LimeHunter7 Mar 27 at 12:18

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.