GeezerVison is showing github.com — sqlobject/sqlobject/util/csvexport.py at master · sqlobject/sqlobject

Expand file tree

csvexport.py

Copy path

More file actions

Latest commit

History

210 lines (176 loc) · 6.88 KB

csvexport.py

Copy path

File metadata and controls

210 lines (176 loc) · 6.88 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

"""

Exports a SQLObject class (possibly annotated) to a CSV file.

"""

import csv

import os

try:

from cStringIO import StringIO

except ImportError:

try:

from StringIO import StringIO

except ImportError:

from io import StringIO, BytesIO

import sqlobject

from sqlobject.compat import PY2, string_type

__all__ = ['export_csv', 'export_csv_zip']

def export_csv(soClass, select=None, writer=None, connection=None,

orderBy=None):

"""

Export the SQLObject class ``soClass`` to a CSV file.

``soClass`` can also be a SelectResults object, as returned by

``.select()``. If it is a class, all objects will be retrieved,

ordered by ``orderBy`` if given, or the ``.csvOrderBy`` attribute

if present (but csvOrderBy will only be applied when no select

result is given).

You can also pass in select results (or simply a list of

instances) in ``select`` -- if you have a list of objects (not a

SelectResults instance, as produced by ``.select()``) then you must

pass it in with ``select`` and pass the class in as the first

argument.

``writer`` is a ``csv.writer()`` object, or a file-like object.

If not given, the string of the file will be returned.

Uses ``connection`` as the data source, if given, otherwise the

default connection.

Columns can be annotated with ``.csvTitle`` attributes, which will

form the attributes of the columns, or 'title' (secondarily), or

if nothing then the column attribute name.

If a column has a ``.noCSV`` attribute which is true, then the

column will be suppressed.

Additionally a class can have an ``.extraCSVColumns`` attribute,

which should be a list of strings/tuples. If a tuple, it should

be like ``(attribute, title)``, otherwise it is the attribute,

which will also be the title. These will be appended to the end

of the CSV file; the attribute will be retrieved from instances.

Also a ``.csvColumnOrder`` attribute can be on the class, which is

the string names of attributes in the order they should be

presented.

"""

return_fileobj = None

if not writer:

return_fileobj = StringIO()

writer = csv.writer(return_fileobj)

elif not hasattr(writer, 'writerow'):

writer = csv.writer(writer)

if isinstance(soClass, sqlobject.SQLObject.SelectResultsClass):

assert select is None, (

"You cannot pass in a select argument (%r) "

"and a SelectResults argument (%r) for soClass" %

(select, soClass))

select = soClass

soClass = select.sourceClass

elif select is None:

select = soClass.select()

if getattr(soClass, 'csvOrderBy', None):

select = select.orderBy(soClass.csvOrderBy)

if orderBy:

select = select.orderBy(orderBy)

if connection:

select = select.connection(connection)

_actually_export_csv(soClass, select, writer)

if return_fileobj:

# They didn't pass any writer or file object in, so we return

# the string result:

return return_fileobj.getvalue()

def _actually_export_csv(soClass, select, writer):

attributes, titles = _find_columns(soClass)

writer.writerow(titles)

for soInstance in select:

row = [getattr(soInstance, attr)

for attr in attributes]

writer.writerow(row)

def _find_columns(soClass):

order = []

attrs = {}

for col in soClass.sqlmeta.columnList:

if getattr(col, 'noCSV', False):

continue

order.append(col.name)

title = col.name

if hasattr(col, 'csvTitle'):

title = col.csvTitle

elif getattr(col, 'title', None) is not None:

title = col.title

attrs[col.name] = title

for attrDesc in getattr(soClass, 'extraCSVColumns', []):

if isinstance(attrDesc, (list, tuple)):

attr, title = attrDesc

else:

attr = title = attrDesc

order.append(attr)

attrs[attr] = title

if hasattr(soClass, 'csvColumnOrder'):

oldOrder = order

order = soClass.csvColumnOrder

for attr in order:

if attr not in oldOrder:

raise KeyError(

"Attribute %r in csvColumnOrder (on class %r) "

"does not exist as a column or in .extraCSVColumns "

"(I have: %r)" % (attr, soClass, oldOrder))

oldOrder.remove(attr)

order.extend(oldOrder)

titles = [attrs[attr] for attr in order]

return order, titles

def export_csv_zip(soClasses, file=None, zip=None, filename_prefix='',

connection=None):

"""

Export several SQLObject classes into a .zip file. Each

item in the ``soClasses`` list may be a SQLObject class,

select result, or ``(soClass, select)`` tuple.

Each file in the zip will be named after the class name (with

``.csv`` appended), or using the filename in the ``.csvFilename``

attribute.

If ``file`` is given, the zip will be written to that. ``file``

may be a string (a filename) or a file-like object. If not given,

a string will be returnd.

If ``zip`` is given, then the files will be written to that zip

file.

All filenames will be prefixed with ``filename_prefix`` (which may

be a directory name, for instance).

"""

import zipfile

close_file_when_finished = False

close_zip_when_finished = True

return_when_finished = False

if file:

if isinstance(file, string_type):

close_file_when_finished = True

file = open(file, 'wb')

elif zip:

close_zip_when_finished = False

else:

return_when_finished = True

if PY2:

file = StringIO()

else:

# zipfile on python3 requires BytesIO

file = BytesIO()

if not zip:

zip = zipfile.ZipFile(file, mode='w')

try:

_actually_export_classes(soClasses, zip, filename_prefix,

connection)

finally:

if close_zip_when_finished:

zip.close()

if close_file_when_finished:

file.close()

if return_when_finished:

return file.getvalue()

def _actually_export_classes(soClasses, zip, filename_prefix,

connection):

for classDesc in soClasses:

if isinstance(classDesc, (tuple, list)):

soClass, select = classDesc

elif isinstance(classDesc, sqlobject.SQLObject.SelectResultsClass):

select = classDesc

soClass = select.sourceClass

else:

soClass = classDesc

select = None

filename = getattr(soClass, 'csvFilename', soClass.__name__)

if not os.path.splitext(filename)[1]:

filename += '.csv'

filename = filename_prefix + filename

zip.writestr(filename,

export_csv(soClass, select, connection=connection))