Working with vector data

from whitebox_workflows import AttributeField, FieldData, FieldDataType, VectorGeometryType, WbEnvironment

wbe = WbEnvironment()

# Let's begin by downloading the Whitebox Workflows 'Southern_Ontario_roads' sample data
wbe.working_directory = wbw.download_sample_data('Southern_Ontario_roads')
print(f'Data have been stored in: {wbe.working_directory}')

# Read in the roads file
roads = wbe.read_vector('roads_utm.shp')

# Let's see some of the properties of this file...
print(f'Vector geometry type: {roads.header.shape_type}') # A VectorGeometryType.PolyLine
print(f'Min X: {roads.header.x_min}')
print(f'Max X: {roads.header.x_max}')
print(f'Min Y: {roads.header.y_min}')
print(f'Max Y: {roads.header.y_max}')
print(f'Number of records: {roads.num_records}')
print(f'Projection: {roads.projection}')

# To retrieve a vector geometry, use the [] syntax.
geom = roads[100]
print(f'Geometry type: {geom.shape_type}')
print(f'Num vertices: {geom.num_points}')
print(f'Num parts: {geom.num_parts}')
print(f'Min X: {geom.x_min}') # There min/max x, y, measure, and z
print(f'First vertex: ({geom.points[0].x}, {geom.points[0].y})')

# What about the file attributes?
print(f'Num records: {roads.attributes.header.num_records}') # Should be same as roads.num_records
print(f'Num fields: {roads.attributes.get_num_fields()}')

# Retrieve and print the attribute fields
att_fields = roads.get_attribute_fields()
for i in range(len(att_fields)):
    print(att_fields[i])

# What's the index of a specific field, by attribute name?
index_of_road_class = roads.get_attribute_field_num("ROAD_CLASS")

# Let's create a new Vector object...

# First, we need to create the attribute fields used for the new vector
out_att_fields = [
    AttributeField("FID", FieldDataType.Int, 6, 0),
    AttributeField("SRC_FID", FieldDataType.Int, 6, 0),
]

# Now create the Vector itself.
out_roads = wbe.new_vector(VectorGeometryType.PolyLine, out_att_fields, proj=roads.projection)

# You can also add a new attribute field after creating the file.
out_roads.add_attribute_field(att_fields[index_of_road_class])

# Now let's filter all the records to find those geometries with a ROAD_CLASS of indicating major highways
old_progress = -1
fid = 1
for i in range(roads.num_records):
    road_class = roads.get_attribute_value(i, 'ROAD_CLASS').get_as_string().lower()
    if 'freeway' in road_class or 'collector' in road_class or 'expressway' in road_class:
        geom = roads[i]
        out_roads.add_record(geom) # Add the record to the output Vector

        # Create the output attribute record...
        att_rec = roads.get_attribute_record(i)
        rec_data = [
            FieldData.new_int(fid), 
            FieldData.new_int(i+1),
            att_rec[index_of_road_class]
        ]
        fid += 1
        # then add it to the table.
        out_roads.add_attribute_record(rec_data, deleted=False)

    # Update the progress after completing each 1% of the records.
    progress = int((i + 1.0) / roads.num_records * 100.0)
    if progress != old_progress:
        old_progress = progress
        print(f'Progress: {progress}%')


# Write the output to file.
print('Writing vector to file...')
wbe.write_vector(out_roads, 'out_roads.shp')