ogrinspect.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. import argparse
  2. from django.contrib.gis import gdal
  3. from django.core.management.base import BaseCommand, CommandError
  4. from django.utils.inspect import get_func_args
  5. class LayerOptionAction(argparse.Action):
  6. """
  7. Custom argparse action for the `ogrinspect` `layer_key` keyword option
  8. which may be an integer or a string.
  9. """
  10. def __call__(self, parser, namespace, value, option_string=None):
  11. try:
  12. setattr(namespace, self.dest, int(value))
  13. except ValueError:
  14. setattr(namespace, self.dest, value)
  15. class ListOptionAction(argparse.Action):
  16. """
  17. Custom argparse action for `ogrinspect` keywords that require
  18. a string list. If the string is 'True'/'true' then the option
  19. value will be a boolean instead.
  20. """
  21. def __call__(self, parser, namespace, value, option_string=None):
  22. if value.lower() == "true":
  23. setattr(namespace, self.dest, True)
  24. else:
  25. setattr(namespace, self.dest, value.split(","))
  26. class Command(BaseCommand):
  27. help = (
  28. "Inspects the given OGR-compatible data source (e.g., a shapefile) and "
  29. "outputs\na GeoDjango model with the given model name. For example:\n"
  30. " ./manage.py ogrinspect zipcode.shp Zipcode"
  31. )
  32. requires_system_checks = []
  33. def add_arguments(self, parser):
  34. parser.add_argument("data_source", help="Path to the data source.")
  35. parser.add_argument("model_name", help="Name of the model to create.")
  36. parser.add_argument(
  37. "--blank",
  38. action=ListOptionAction,
  39. default=False,
  40. help="Use a comma separated list of OGR field names to add "
  41. "the `blank=True` option to the field definition. Set to `true` "
  42. "to apply to all applicable fields.",
  43. )
  44. parser.add_argument(
  45. "--decimal",
  46. action=ListOptionAction,
  47. default=False,
  48. help="Use a comma separated list of OGR float fields to "
  49. "generate `DecimalField` instead of the default "
  50. "`FloatField`. Set to `true` to apply to all OGR float fields.",
  51. )
  52. parser.add_argument(
  53. "--geom-name",
  54. default="geom",
  55. help="Specifies the model name for the Geometry Field (defaults to `geom`)",
  56. )
  57. parser.add_argument(
  58. "--layer",
  59. dest="layer_key",
  60. action=LayerOptionAction,
  61. default=0,
  62. help="The key for specifying which layer in the OGR data "
  63. "source to use. Defaults to 0 (the first layer). May be "
  64. "an integer or a string identifier for the layer.",
  65. )
  66. parser.add_argument(
  67. "--multi-geom",
  68. action="store_true",
  69. help="Treat the geometry in the data source as a geometry collection.",
  70. )
  71. parser.add_argument(
  72. "--name-field",
  73. help="Specifies a field name to return for the __str__() method.",
  74. )
  75. parser.add_argument(
  76. "--no-imports",
  77. action="store_false",
  78. dest="imports",
  79. help="Do not include `from django.contrib.gis.db import models` statement.",
  80. )
  81. parser.add_argument(
  82. "--null",
  83. action=ListOptionAction,
  84. default=False,
  85. help="Use a comma separated list of OGR field names to add "
  86. "the `null=True` option to the field definition. Set to `true` "
  87. "to apply to all applicable fields.",
  88. )
  89. parser.add_argument(
  90. "--srid",
  91. help="The SRID to use for the Geometry Field. If it can be "
  92. "determined, the SRID of the data source is used.",
  93. )
  94. parser.add_argument(
  95. "--mapping",
  96. action="store_true",
  97. help="Generate mapping dictionary for use with `LayerMapping`.",
  98. )
  99. def handle(self, *args, **options):
  100. data_source, model_name = options.pop("data_source"), options.pop("model_name")
  101. # Getting the OGR DataSource from the string parameter.
  102. try:
  103. ds = gdal.DataSource(data_source)
  104. except gdal.GDALException as msg:
  105. raise CommandError(msg)
  106. # Returning the output of ogrinspect with the given arguments
  107. # and options.
  108. from django.contrib.gis.utils.ogrinspect import _ogrinspect, mapping
  109. # Filter options to params accepted by `_ogrinspect`
  110. ogr_options = {
  111. k: v
  112. for k, v in options.items()
  113. if k in get_func_args(_ogrinspect) and v is not None
  114. }
  115. output = [s for s in _ogrinspect(ds, model_name, **ogr_options)]
  116. if options["mapping"]:
  117. # Constructing the keyword arguments for `mapping`, and
  118. # calling it on the data source.
  119. kwargs = {
  120. "geom_name": options["geom_name"],
  121. "layer_key": options["layer_key"],
  122. "multi_geom": options["multi_geom"],
  123. }
  124. mapping_dict = mapping(ds, **kwargs)
  125. # This extra legwork is so that the dictionary definition comes
  126. # out in the same order as the fields in the model definition.
  127. rev_mapping = {v: k for k, v in mapping_dict.items()}
  128. output.extend(
  129. [
  130. "",
  131. "",
  132. "# Auto-generated `LayerMapping` dictionary for %s model"
  133. % model_name,
  134. "%s_mapping = {" % model_name.lower(),
  135. ]
  136. )
  137. output.extend(
  138. " '%s': '%s'," % (rev_mapping[ogr_fld], ogr_fld)
  139. for ogr_fld in ds[options["layer_key"]].fields
  140. )
  141. output.extend(
  142. [
  143. " '%s': '%s',"
  144. % (options["geom_name"], mapping_dict[options["geom_name"]]),
  145. "}",
  146. ]
  147. )
  148. return "\n".join(output)