Question

如果我有一个名为list_files的函数（将目录作为参数，并通过文件夹结构进行“ os.walks”），而另一个名为main的函数（将文件作为参数并将其保存到DB），我将如何嵌套它们，以便目录结构中的每个文件都写入数据库？我在下面尝试了for循环，但无法正常工作。

我知道os.walk返回一个需要转换为字符串的列表。而且我知道我做错了。

"""A tool for saving files to and from a postgresql db.
"""
import os
import sys
import argparse
import psycopg2

db_conn_str = "postgresql://word:word@111.11.111.1:5432/DBNAME"
create_table_stm = """
CREATE TABLE IF NOT EXISTS files (
    id serial primary key,
    orig_filename text not null,
    file_data bytea not null
)
"""

# Walk through the directory
def list_files(dir):
    r = []
    for root, dirs, files in os.walk(dir):
        for name in files:
            r.append(os.path.join(root, name))
    return r


def main(argv):
    for each_file in list_files(''.join(r)):
        parser = argparse.ArgumentParser()
        parser_action = parser.add_mutually_exclusive_group(required=True)
        parser_action.add_argument("--store", action='store_const', const=True, help="Load an image from the named file and save it in the DB")
        parser_action.add_argument("--fetch", type=int, help="Fetch an image from the DB and store it in the named file, overwriting it if it exists. Takes the database file identifier as an argument.", metavar='42')
        parser.add_argument("filename", help="Name of file to write to / fetch from")
        args = parser.parse_args(argv[1:])

        conn = psycopg2.connect(db_conn_str)
        curs = conn.cursor()


        if args.store:
                # This should just be f.read() (dropping the psycopg2.Binary) when using Python 3 -- the current version is written for Python 2.7
            with open(args.filename,'rb') as f:
                filedata = psycopg2.Binary(f.read())
                curs.execute("INSERT INTO files(id, orig_filename, file_data) VALUES (DEFAULT,%s,%s) RETURNING id", (args.filename, filedata))
                print curs
                returned_id = curs.fetchone()[0]
            print("Stored {0} into DB record {1}".format(args.filename, returned_id))
            conn.commit()

        elif args.fetch is not None:
            with open(args.filename,'wb') as f:
                curs.execute("SELECT file_data, orig_filename FROM files WHERE id = %s", (int(args.fetch),))
                (file_data, orig_filename) = curs.fetchone()
                f.write(file_data)
            print("Fetched {0} into file {1}; original filename was {2}".format(args.fetch, args.filename, orig_filename))

        conn.close()


if __name__ == '__main__':
    main(sys.argv)

结合两个功能：遍历目录，将每个文件写入DB

0 个答案: