[Tutor] How to invoke different member-function according to different object?

[Peter Otten]

daedae11 wrote:

> I was asked to write a program to move files between ZIP(.zip) and
> TAR/GZIP(.tgz/.tar.gz) or TAR/BZIP2(.tbz/.tar.bz2) archive.
> 
> my code is:
> 
> 
> import zipfile;
> import tarfile;
> import os;
> from os import path ;
> 
> def showAllFiles(fileObj):
>     if fileObj.filename.endswith("zip"):
>         if isinstance(fileObj, zipfile.ZipFile):
>             print "j"*20;
>         for name in fileObj.namelist():
>             print name;
>     else:
>         for name in fileObj.getnames():
>             print name;
> 
> def moveFile(srcObj, dstObj):
>     fileName = raw_input("input the name of the file to move: ");
>     srcObj.extract(fileName);
>     if isinstance(dstObj, zipfile.ZipFile):
>         dstObj.write(fileName);
>     else:
>         dstObj.addfile(tarfile.TarInfo(fileName));
>     os.remove(fileName);
>     
> def main():
>     intro = """
> enter a choice
> (M)ove file from source file to destinatiom file
> (S)how all the files in source file
> (Q)uit
> your choice is: """
>     srcFile = raw_input("input the source file name: ");
>     dstFile = raw_input("input the destination file name: ");
>     while True:
>         with ( zipfile.ZipFile(srcFile, "r") if srcFile.endswith("zip")
>         else tarfile.open(srcFile, "r"+":"+path.splitext(srcFile)[1][1:])
>         ) as srcObj, \ ( zipfile.ZipFile(dstFile, "r") if
>            dstFile.endswith("zip") else
>             tarfile.open(dstFile, "w"+":"+path.splitext(dstFile)[1][1:]) )
>             as dstObj:
>                 choice = raw_input(intro)[0].lower();
>                 if choice == "s":
>                     showAllFiles(srcObj);
>                 elif choice == "m":
>                     moveFile(srcObj, dstObj);
>                 elif choice == "q":
>                     break;
>                 else:
>                     print "invalid command!"
> 
> if __name__ == '__main__':
>     main();
> 
> But there are some problems.
> 1. It could extract file successfully, but can't add files to .tar.gz
> file. 2. I think it's a little tedious, but I don't know how to improve
> it.
> 
> Please  give me some help , thank you!

First, but least: don't pollute your code with trailing semicola.
Second: what Steven says, plus it will always take you longer to implement a 
script than you initially think. I ran into that for the umpteenth time when 
trying to reorganize your script.

Now to your code:
- Try to separate the user interface (aka raw_input()) from the other 
functionality of your script. moveFile() will be easier to test when you  
pass the filename as an argument rather than asking the user for it. 

- Don't overuse inlined if-else expressions. A separate

if complex condition:
   var = ...
else:
   var = ...

is easier to read -- and easier to extend with another elif.

- If you have a lot of if-else switches consider classes

# wrong
animal = dog_or_duck()
if its_a_dog(animal):
   bark()
elif its_a_duck(animal):
   quack()

# right
animal = dog_or_duck()
animal.make_some_noise()

In your case there already are classes, but with incompatible interfaces. 
You can solve that by subclassing or making wrapper classes.

Another approach that addresses the same problem is table-driven:

def quack(): print "quack"
def bark(): print "wuff"
lookup_noise = {
    Dog: bark,
    Duck: quack,
}
animal = dog_or_duck()
lookup_noise[animal.__class__]()

In my rewritten version of your script I use both approaches. It is still 
incomplete and has already become almost too long to post. It is probably a 
bit too complex for a beginner, too, but the idea is simple and if you have 
questions I'm willing to answer. 

Basically there is a class for every archive type which has to implement an 
open(name) method which has to return a file-like object and an add(name, 
file) which has to add a file-like object as name to the archive.
You can ask if a certain archive class can deal with a file with e. g.

ZipArchive.can_handle(filename)

The script features an adhoc test that you can invoke with

python scriptname.py --test

Look into http://docs.python.org/library/unittest.html if you are serious 
about testing. I think you should be ;)

import os
import shutil
import sys
import tarfile
import zipfile

def _not_implemented(obj, methodname):
    typename = obj.__class__.__name__
    print >> sys.stderr, "{}.{}() not implemented".format(typename, 
methodname)

class Archive(object):
    @classmethod
    def can_handle(class_, filename):
        return filename.lower().endswith(class_.suffixes)

    def __init__(self, filename, mode):
        if mode not in ("r", "w"):
            raise ValueError("Mode {} not understood".format(mode))
        self.filename = filename
        self.mode = mode
        self._archive = self._open_archive()
    def getnames(self):
        return self._archive.getnames()
    def __enter__(self):
        return self
    def __exit__(self, *args):
        self._archive.close()

    def _open_archive(self):
        _not_implemented(self, "_open")
    def add(self, name, srcfile):
        _not_implemented(self, "add")
    def open(self, name):
        _not_implemented(self, "extract")

class ZipArchive(Archive):
    suffixes = (".zip",)
    def _open_archive(self):
        return zipfile.ZipFile(self.filename, self.mode)
    def getnames(self):
        return self._archive.namelist()
    def open(self, name):
        return self._archive.open(name)

class TarArchive(Archive):
    suffixes = (".tar",)
    def _open_archive(self):
        return tarfile.open(self.filename, self.mode)
    def add(self, name, srcfile):
        tmpname = "tmp.xxx" # TODO use tempfile module
        with open(tmpname, "wb") as dstfile:
            shutil.copyfileobj(srcfile, dstfile)
        self._archive.add(tmpname, name)
        os.remove(tmpname)

class TgzArchive(TarArchive):
    suffixes = (".tar.gz", ".tgz")
    def _open_archive(self):
        return tarfile.open(self.filename, self.mode + ":gz")

lookup = [
    ZipArchive,
    TarArchive,
    TgzArchive,
]

def move_file(src, dst, filename):
    dst.add(filename, src.open(filename))

def open_archive(filename, mode):
    for Archive in lookup:
        if Archive.can_handle(filename):
            return Archive(filename, mode)
    raise ValueError("Archive type of {!r} not understood".format(filename))

def show_all_files(archive):
    for name in archive.getnames():
        print name

def main():
    if "--test" in sys.argv:
        with zipfile.ZipFile("one.zip", "w") as dst:
            for data in "alpha beta gamma".split():
                dst.writestr(data + ".txt", data)
        with open_archive("one.zip", "r") as src:
            with open_archive("two.tar", "w") as dst:
                move_file(src, dst, "alpha.txt")
            with open_archive("three.tar.gz", "w") as dst:
                move_file(src, dst, "beta.txt")
                move_file(src, dst, "gamma.txt")
        with tarfile.open("two.tar") as src:
            assert src.getnames() == ["alpha.txt"]
        with tarfile.open("three.tar.gz") as src:
            assert src.getnames() == ["beta.txt", "gamma.txt"]
        # TODO: check contents
        return

    intro = """
enter a choice
(M)ove file from source file to destinatiom file
(S)how all the files in source file
(Q)uit
your choice is: """
    srcname = raw_input("input the source file name: ")
    dstname = raw_input("input the destination file name: ")
    with open_archive(srcname, "r") as src:
        while True:
            choice = raw_input(intro).strip().lower()
            if choice == "s":
                show_all_files(src)
            elif choice == "m":
                filename = raw_input("input the name of the file to move: ")
                with open_archive(dstname, "w") as dst:
                    move_file(src, dst, filename)
            elif choice == "q":
                break
            else:
                print "invalid command ({!r})!" % choice

if __name__ == '__main__':
    main()