Skip to content

Instantly share code, notes, and snippets.

@rik
Created October 14, 2022 06:36
Show Gist options
  • Save rik/c15479068812eef069951270e71e3e91 to your computer and use it in GitHub Desktop.
Save rik/c15479068812eef069951270e71e3e91 to your computer and use it in GitHub Desktop.
import argparse
import csv
from collections.abc import Generator
from contextlib import contextmanager, closing
from io import TextIOWrapper
from typing import Any
from django.core.management.base import BaseCommand
from django.db import transaction
from example.core.models import Author, Book
class Command(BaseCommand):
help = "Import books from a CSV"
def add_arguments(self, parser):
parser.add_argument(
"--write",
action="store_true",
default=False,
help="Actually edit the database",
)
parser.add_argument("file", type=argparse.FileType())
def handle(
self, *args: Any, file: TextIOWrapper, write: bool, **kwargs: Any
) -> None:
if not write:
self.stdout.write("In dry run mode (--write not passed)")
books_created = 0
authors_created = 0
with closing(file), atomic(commit=write):
csvfile = csv.reader(file)
header = next(csvfile)
if header != ["Title", "Author"]:
self.stdout.write("Unexpected header row, should be “Title,Author”")
raise SystemExit(1)
for title, author_name in csvfile:
author, created = Author.objects.get_or_create(name=author_name)
if created:
authors_created += 1
book, created = Book.objects.get_or_create(
title=title,
author=author,
)
if created:
books_created += 1
if write:
prefix = "Created"
else:
prefix = "Would create"
self.stdout.write(
f"{prefix} {books_created} books and {authors_created} authors."
)
class DoRollback(Exception):
pass
@contextmanager
def atomic(*, commit: bool) -> Generator[None, None, None]:
try:
with transaction.atomic():
yield
if not commit:
raise DoRollback()
except DoRollback:
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment