* python/tag.cc, tests/test_tagfile.py:

  - add support a filename argument in apt_pkg.TagFile() (in addition
    to the file object currently supported)

3 files changed, 53 insertions, 24 deletions

diff --git a/debian/changelog b/debian/changelog
index 89c07119..5813656a 100644
--- a/debian/changelog
+++ b/debian/changelog
@@ -10,6 +10,9 @@ python-apt (0.8.4) UNRELEASED; urgency=low
   * tests/test_tagfile.py:
     - add test for apt_pkg.TagFile() both for compressed/uncompressed
       files
+  * python/tag.cc, tests/test_tagfile.py:
+    - add support a filename argument in apt_pkg.TagFile() (in addition
+      to the file object currently supported)
 
  -- Michael Vogt <michael.vogt@ubuntu.com>  Wed, 04 Jan 2012 12:07:48 +0100
 
diff --git a/python/tag.cc b/python/tag.cc
index 3996af8e..b680dc02 100644
--- a/python/tag.cc
+++ b/python/tag.cc
@@ -390,24 +390,50 @@ PyObject *ParseSection(PyObject *self,PyObject *Args)
 
 static PyObject *TagFileNew(PyTypeObject *type,PyObject *Args,PyObject *kwds)
 {
+   TagFileData *New;
    PyObject *File;
+
    char *kwlist[] = {"file", 0};
    if (PyArg_ParseTupleAndKeywords(Args,kwds,"O",kwlist,&File) == 0)
       return 0;
-   int fileno = PyObject_AsFileDescriptor(File);
-   if (fileno == -1)
+
+   // check if we got a filename or a file object
+   int fileno = -1;
+   const char *filename = NULL;
+   if (PyString_Check(File))
+      filename = PyObject_AsString(File);
+   else
+      fileno = PyObject_AsFileDescriptor(File);
+
+   // handle invalid arguments
+   if (fileno == -1 && filename == NULL)
+   {
+      PyErr_SetString(PyExc_TypeError, 
+                      "Argument must be string, fd or have a fileno() method");
       return 0;
+   }
 
-   TagFileData *New = (TagFileData*)type->tp_alloc(type, 0);
+   if (fileno > 0)
+   {
+      New = (TagFileData*)type->tp_alloc(type, 0);
 #ifdef APT_HAS_GZIP
-   new (&New->Fd) FileFd();
-   New->Fd.OpenDescriptor(fileno, FileFd::ReadOnlyGzip, false);
+      new (&New->Fd) FileFd();
+      New->Fd.OpenDescriptor(fileno, FileFd::ReadOnlyGzip, false);
 #else
-   new (&New->Fd) FileFd(fileno,false);
+      new (&New->Fd) FileFd(fileno,false);
 #endif
-   New->Owner = File;
-   Py_INCREF(New->Owner);
-   new (&New->Object) pkgTagFile(&New->Fd);
+      New->Owner = File;
+      Py_INCREF(New->Owner);
+      new (&New->Object) pkgTagFile(&New->Fd);
+   }
+   else 
+   {
+      New = (TagFileData*)type->tp_alloc(type, 0);
+      new (&New->Fd) FileFd(filename, FileFd::ReadOnly, FileFd::Extension, false);
+      New->Owner = File;
+      Py_INCREF(New->Owner);
+      new (&New->Object) pkgTagFile(&New->Fd);
+   } 
 
    // Create the section
    New->Section = (TagSecData*)(&PyTagSection_Type)->tp_alloc(&PyTagSection_Type, 0);
diff --git a/tests/test_tagfile.py b/tests/test_tagfile.py
index b14dd9b4..371cc6ee 100644
--- a/tests/test_tagfile.py
+++ b/tests/test_tagfile.py
@@ -7,6 +7,7 @@
 # notice and this notice are preserved.
 """Unit tests for verifying the correctness of apt_pkg.TagFile"""
 
+import glob
 import os
 import unittest
 
@@ -21,22 +22,21 @@ class TestTagFile(unittest.TestCase):
 
     def test_tag_file(self):
         basepath = os.path.dirname(__file__)
-        tagfilepath = os.path.join(
-            basepath, "./data/tagfile/history.log")
-        tagfile = apt_pkg.TagFile(open(tagfilepath))
-        for i, stanza in enumerate(tagfile):
-            pass
-        self.assertEqual(i, 2)
+        tagfilepath = os.path.join(basepath, "./data/tagfile/*")
+        # test once for compressed and uncompressed
+        for testfile in glob.glob(tagfilepath):
+            # test once using the open() method and once using the path
+            for f in [testfile, open(testfile)]:
+                tagfile = apt_pkg.TagFile(f)
+                for i, stanza in enumerate(tagfile):
+                    pass
+                self.assertEqual(i, 2)
 
-    def test_tag_file_compressed(self):
-        basepath = os.path.dirname(__file__)
-        tagfilepath = os.path.join(
-            basepath, "./data/tagfile/history.1.log.gz")
-        tagfile = apt_pkg.TagFile(open(tagfilepath))
-        for i, stanza in enumerate(tagfile):
-            #print stanza
-            pass
-        self.assertEqual(i, 2)
+    def test_errors(self):
+        # Raises SystemError via lbiapt
+        self.assertRaises(SystemError, apt_pkg.TagFile, "not-there-no-no")
+        # Raises Type error
+        self.assertRaises(TypeError, apt_pkg.TagFile, object())
 
 if __name__ == "__main__":
     unittest.main()