diff options
Diffstat (limited to 'lib')
-rw-r--r-- | lib/GTFS.hs | 70 | ||||
-rw-r--r-- | lib/Persist.hs | 8 |
2 files changed, 64 insertions, 14 deletions
diff --git a/lib/GTFS.hs b/lib/GTFS.hs index be80745..bd29b6d 100644 --- a/lib/GTFS.hs +++ b/lib/GTFS.hs @@ -6,6 +6,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE StandaloneKindSignatures #-} @@ -20,6 +21,7 @@ module GTFS where import qualified Codec.Archive.Zip as Zip +import Control.Monad.ST (runST) import Data.Aeson (FromJSON, Options (fieldLabelModifier), ToJSON, defaultOptions, @@ -30,6 +32,7 @@ import qualified Data.ByteString as BS import qualified Data.ByteString.Lazy as LB import Data.Csv ((.:)) import qualified Data.Csv as CSV +import Data.Function (on) import Data.Functor ((<&>)) import Data.Kind (Type) import Data.Maybe (fromJust, fromMaybe) @@ -41,6 +44,7 @@ import Data.Time.Calendar.MonthDay (monthAndDayToDayOfYearValid) import qualified Data.Time.Calendar.OrdinalDate as Day import Data.Vector (Vector) import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Intro as V import Fmt ((+|), (|+)) import GHC.Generics (Generic) import Text.Regex.TDFA ((=~)) @@ -195,16 +199,16 @@ data Trip (deep :: Depth) = Trip , tripServiceId :: Text -- , tripWheelchairAccessible :: Bool -- , tripBikesAllowed :: Bool - , tripShapeId :: Text + , tripShape :: Switch deep Shape Text , tripStops :: Optional deep (Vector (Stop deep)) } deriving Generic deriving instance Show (Trip Shallow) deriving instance Show (Trip Deep) -instance FromJSON (Optional d (Vector (Stop d))) => FromJSON (Trip d) where +instance (FromJSON (Switch d Shape Text), FromJSON (Optional d (Vector (Stop d)))) => FromJSON (Trip d) where parseJSON = genericParseJSON (aesonOptions "trip") -instance ToJSON (Optional d (Vector (Stop d))) => ToJSON (Trip d) where +instance (ToJSON (Switch d Shape Text), ToJSON (Optional d (Vector (Stop d)))) => ToJSON (Trip d) where toJSON = genericToJSON (aesonOptions "trip") instance ToSchema (Trip Deep) where declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "trip") @@ -213,6 +217,25 @@ instance ToSchema (Trip Deep) where tableLookup :: Eq key => (a -> key) -> key -> Vector a -> Maybe a tableLookup proj key = V.find (\a -> proj a == key) +data ShapePoint = ShapePoint + { shapePtId :: Text + , shapePtLat :: Double + , shapePtLong :: Double + , shapePtSequence :: Int + } deriving Generic + +data Shape = Shape + { shapeId :: Text + , shapePoints :: Vector (Double,Double) + } deriving (Generic, Show) + +instance FromJSON Shape where + parseJSON = genericParseJSON (aesonOptions "shape") +instance ToJSON Shape where + toJSON = genericToJSON (aesonOptions "shape") +instance ToSchema Shape where + declareNamedSchema = genericDeclareNamedSchema (swaggerOptions "shape") + instance CSV.FromNamedRecord Station where parseNamedRecord r = Station <$> r .: "stop_id" @@ -241,6 +264,13 @@ instance CSV.FromNamedRecord Calendar where <*> r .: "start_date" <*> r .: "end_date" +instance CSV.FromNamedRecord ShapePoint where + parseNamedRecord r = ShapePoint + <$> r .: "shape_id" + <*> r .: "shape_pt_lat" + <*> r .: "shape_pt_lon" + <*> r .: "shape_pt_sequence" + intAsBool :: CSV.NamedRecord -> BS.ByteString -> CSV.Parser (Maybe Bool) intAsBool r field = do int <- r .: field @@ -287,6 +317,7 @@ data RawGTFS = RawGTFS , rawTrips :: Vector (Trip Shallow) , rawCalendar :: Maybe (Vector Calendar) , rawCalendarDates :: Maybe (Vector CalendarDate) + , rawShapePoints :: Maybe (Vector ShapePoint) } @@ -295,6 +326,7 @@ data GTFS = GTFS , trips :: Map TripID (Trip Deep) , calendar :: Map DayOfWeek (Vector Calendar) , calendarDates :: Map Day (Vector CalendarDate) + , shapes :: Map Text Shape , fancyCalendar :: Day -> (Vector ServiceID, Vector (Trip Deep)) -- ^ a more "fancy" encoding of the calendar? @@ -311,6 +343,7 @@ loadRawGtfs path = do <*> decodeTable' "trips.txt" zip <*> decodeTable "calendar.txt" zip <*> decodeTable "calendar_dates.txt" zip + <*> decodeTable "shapes.txt" zip where decodeTable :: CSV.FromNamedRecord a => FilePath -> Zip.Archive -> IO (Maybe (Vector a)) decodeTable path zip = @@ -327,8 +360,13 @@ loadRawGtfs path = do loadGtfs :: FilePath -> IO GTFS loadGtfs path = do shallow@RawGTFS{..} <- loadRawGtfs path + -- TODO: sort these according to sequence numbers + let shapes = + V.foldr' sortShapePoint mempty + $ V.modify (V.sortBy (compare `on` shapePtSequence)) + (fromMaybe mempty rawShapePoints) stops' <- V.mapM (pushStop rawStations) rawStops - trips' <- V.mapM (pushTrip stops') rawTrips + trips' <- V.mapM (pushTrip stops' shapes) rawTrips pure $ GTFS { stations = M.fromList $ (\station -> (stationId station, station)) @@ -345,6 +383,7 @@ loadGtfs path = do fmap V.fromList $ M.fromListWith (<>) $ (\cd -> (caldateDate cd, [cd])) <$> V.toList (fromMaybe mempty rawCalendarDates) + , shapes } where weekdays Calendar{..} = @@ -361,13 +400,24 @@ loadGtfs path = do Just a -> pure a Nothing -> fail $ "station with id "+|stopStation stop|+"is mentioned but not defined." pure $ stop { stopStation = station } - pushTrip :: Vector (Stop Deep) -> Trip Shallow -> IO (Trip Deep) - pushTrip stops trip = if V.length alongRoute < 2 + pushTrip :: Vector (Stop Deep) -> Map Text Shape -> Trip Shallow -> IO (Trip Deep) + pushTrip stops shapes trip = if V.length alongRoute < 2 then fail $ "trip with id "+|tripTripID trip|+" has no stops" - else pure $ trip { tripStops = alongRoute } - where alongRoute = -- TODO: sort these according to stops - V.filter (\s -> stopTrip s == tripTripID trip) stops - + else do + a <- case M.lookup (tripShape trip) shapes of + Nothing -> fail $ "trip with id "+|tripTripID trip|+" mentions a shape that does not exist." + Just a -> pure a + pure $ trip { tripStops = alongRoute, tripShape = a } + where alongRoute = + V.modify (V.sortBy (compare `on` stopSequence)) + $ V.filter (\s -> stopTrip s == tripTripID trip) stops + sortShapePoint :: ShapePoint -> Map Text Shape -> Map Text Shape + sortShapePoint ShapePoint{..} shapes = M.alter appendPoint shapePtId shapes + where + point = (shapePtLat, shapePtLong) + appendPoint = \case + Just shape -> Just $ shape { shapePoints = V.cons point (shapePoints shape) } + Nothing -> Just $ Shape { shapeId = shapePtId, shapePoints = V.singleton point } servicesOnDay :: GTFS -> Day -> Vector ServiceID diff --git a/lib/Persist.hs b/lib/Persist.hs index 9486a1d..4a6d9b4 100644 --- a/lib/Persist.hs +++ b/lib/Persist.hs @@ -37,15 +37,15 @@ import Control.Monad.Logger (NoLoggingT) import Control.Monad.Reader (ReaderT) import Data.Data (Proxy (..)) import Data.Pool (Pool) -import Data.Time (NominalDiffTime, +import Data.Time (NominalDiffTime, TimeOfDay, UTCTime (utctDay), addUTCTime, dayOfWeek, diffUTCTime, - getCurrentTime, nominalDay, TimeOfDay) + getCurrentTime, nominalDay) +import Data.Time.Calendar (Day, DayOfWeek (..)) +import Data.Vector (Vector) import Database.Persist.Postgresql (SqlBackend) import GHC.Generics (Generic) import Web.PathPieces (PathPiece) -import Data.Vector (Vector) -import Data.Time.Calendar (Day, DayOfWeek (..)) newtype Token = Token UUID deriving newtype |