-- | -- Copyright : Anders Claesson 2012-2014 -- Maintainer : Anders Claesson -- License : BSD-3 -- import Data.Text (Text) import qualified Data.Text as T import qualified Data.Text.IO as IO import Data.Text.Encoding (decodeUtf8) import Network.HTTP (urlEncodeVars) import Network.Curl.Download (openURI) import Options.Applicative import Sloane.Config import Sloane.DB hiding (null) import qualified Sloane.DB as DB type URL = String type Seq = Text data Visibility = Visible | Internal data Options = Cmd Command | IntOpts SearchOpts -- Internal opts for fallback to lookup data Command = Lookup SearchOpts | Grep SearchOpts | Filter FilterOpts | Update | Version data SearchOpts = SearchOpts { full :: Bool -- Print all fields? , keys :: String -- Keys of fields to print , limit :: Int -- Fetch at most this many entries , url :: Bool -- Print URLs of found entries , terms :: [String] -- Search terms } data FilterOpts = FilterOpts { invert :: Bool -- Return sequences NOT in DB } oeisKeys :: String oeisKeys = "ISTUVWXNDHFYAOEeptoKC" -- Valid OEIS keys oeisUrls :: Config -> DB -> [URL] oeisUrls cfg = map ((oeisHost cfg ++) . T.unpack) . aNumbers oeisLookup :: SearchOpts -> Config -> IO DB oeisLookup opts cfg = (parseOEISEntries . decodeUtf8 . either error id) <$> openURI (oeisURL cfg ++ "&" ++ urlEncodeVars [("n", show n), ("q", q)]) where n = limit opts q = unwords $ terms opts grepDB :: SearchOpts -> DB -> DB grepDB opts = DB.take n . DB.grep (T.pack q) where n = limit opts q = unwords $ terms opts filterDB :: FilterOpts -> DB -> IO [Seq] filterDB opts db = filter match . parseSeqs <$> IO.getContents where match q = (if invert opts then id else not) (DB.null $ DB.grep q db) parseSeqs = filter (not . T.null) . map mkSeq . T.lines mkSeq = normalize . dropComment dropComment = T.takeWhile (/= '#') normalize = T.intercalate (T.pack ",") . T.words . clean . T.map tr tr c = if c `elem` ";," then ' ' else c clean = T.filter (`elem` " 0123456789-") searchOptionsParser :: Visibility -> Parser SearchOpts searchOptionsParser visibility = hiddenHelp <*> (SearchOpts <$> switch ( short 'a' <> long "all" <> help "Print all fields" <> f ) <*> strOption ( short 'k' <> metavar "KEYS" <> value "SN" <> help "Keys of fields to print [default: SN]" <> f ) <*> option auto ( short 'n' <> metavar "N" <> value 5 <> help "Fetch at most this many entries [default: 5]" <> f ) <*> switch ( long "url" <> help "Print URLs of found entries" <> f ) <*> some (argument str (metavar "TERMS..."))) where f = case visibility of {Visible -> idm; Internal -> internal} filterOptionsParser :: Parser FilterOpts filterOptionsParser = FilterOpts <$> switch (long "invert" <> help "Return sequences NOT in the database") commandParser :: Parser Command commandParser = subparser ( command "lookup" (info (Lookup <$> searchOptionsParser Visible) ( progDesc "Lookup a sequence, or other search term, in OEIS" )) <> command "grep" (info (Grep <$> searchOptionsParser Visible) ( progDesc "Grep for a sequence in the local database" )) <> command "filter" (info (Filter <$> filterOptionsParser) ( progDesc ("Read sequences from stdin and " ++ "return those that are in the local database"))) <> command "update" (info (pure Update) ( progDesc "Update the local database" )) <> command "version" (info (pure Version) ( progDesc "Show version info" )) ) optionsParser :: Parser Options optionsParser = (Cmd <$> commandParser) <|> (IntOpts <$> searchOptionsParser Internal) runSearch :: (SearchOpts -> Config -> IO DB) -> SearchOpts -> Config -> IO () runSearch f opts cfg = f opts cfg >>= if url opts then putStr . unlines . oeisUrls cfg else putDB cfg (if full opts then oeisKeys else keys opts) runCmd :: Command -> Config -> IO () runCmd (Lookup opts) = runSearch oeisLookup opts runCmd (Grep opts) = runSearch (\o cfg -> grepDB o <$> readDB cfg) opts runCmd (Filter opts) = \c -> readDB c >>= filterDB opts >>= mapM_ IO.putStrLn runCmd Update = initDB runCmd Version = putStrLn . name hiddenHelp :: Parser (a -> a) hiddenHelp = abortOption ShowHelpText $ hidden <> short 'h' <> long "help" main :: IO () main = do conf <- defaultConfig opts <- customExecParser preferences (info parser description) case opts of (Cmd cmd) -> runCmd cmd conf (IntOpts o) -> runCmd (Lookup o) conf -- Fallback to 'lookup' where parser = hiddenHelp <*> optionsParser preferences = prefs showHelpOnError description = fullDesc <> footer "Run 'sloane COMMAND --help' for help on a specific command."