sloane-1.8.2: sloane.hs
-- |
-- Copyright : Anders Claesson 2012-2014
-- Maintainer : Anders Claesson <anders.claesson@gmail.com>
-- License : BSD-3
--
import qualified Codec.Compression.GZip as GZip
import Control.Monad (unless, when, liftM2)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as IO
import Data.Text.Encoding (decodeUtf8)
import Data.Maybe (fromJust)
import Data.Time (diffUTCTime, getCurrentTime)
import Network.HTTP
import Network.URI (parseURI)
import Options.Applicative
import System.Console.ANSI
import System.Console.Terminal.Size (Window (..), size)
import System.Directory
import System.FilePath ((</>))
import System.IO (stderr)
type OEISEntries = [String]
type ANumbers = [String]
type Query = String
type Keys = String
data Args = Args
Bool -- a: Print all fields?
String -- k: Keys of fields to print
Int -- n: Fetch at most this many entries
Bool -- invert: Return sequences NOT in OEIS
Bool -- update: Update the local sequence cache
Bool -- url: Print URLs of found entries
Bool -- version
[String] -- search terms
name, oeisHost, oeisURL, cacheDir, cacheFile, cacheURL :: String
name = "sloane 1.8.2"
oeisHost = "http://oeis.org/"
oeisURL = oeisHost ++ "search?fmt=text"
cacheDir = ".sloane"
cacheFile = "stripped.gz"
cacheURL = oeisHost ++ cacheFile
msgDownloadingCache, msgCacheIsUpToDate, msgNoCache, msgOldCache :: String
msgDownloadingCache = unlines
[ "Downloading " ++ cacheURL
, "This may take a minute or two ..."
]
msgCacheIsUpToDate = unlines
[ "The sequence cache is now up-to-date"
]
msgNoCache = unlines
[ "No sequence cache found. You need to run \"sloane --update\""
]
msgOldCache = unlines
[ "The sequence cache is more than 100 days old"
, "You may want to run \"sloane --update\""
]
select :: Keys -> OEISEntries -> OEISEntries
select ks = filter (\line -> null line || head line `elem` ks)
aNumbers :: OEISEntries -> ANumbers
aNumbers es = [ words ids !! 1 | ids@(_:_) <- select "I" es ]
urls :: OEISEntries -> String
urls = unlines . map (oeisHost ++ ) . aNumbers
get :: HStream b => String -> IO b
get uri = simpleHTTP (defaultGETRequest_ uri') >>= getResponseBody
where
uri' = fromJust $ parseURI uri
searchOEIS :: Int -> Query -> IO OEISEntries
searchOEIS n s = trim `fmap` get uri
where
trim = map (drop 1) . reverse . drop 2 . reverse . drop 5 . lines
uri = oeisURL ++ "&" ++ urlEncodeVars [("n", show n), ("q", s)]
cropStr :: (Int -> String -> String) -> Int -> String -> String
cropStr f maxLen s = if maxLen < length s then f maxLen s else s
cropSeq :: Int -> String -> String
cropSeq = cropStr $ \maxLen ->
reverse . dropWhile (/= ',') . reverse . take maxLen
cropLine :: Int -> String -> String
cropLine = cropStr $ \maxLen s -> take (maxLen-2) s ++ ".."
getWidth :: IO Int
getWidth = maybe maxBound width `fmap` size
put :: String -> IO ()
put = IO.putStr . T.pack
putErr :: String -> IO ()
putErr = IO.hPutStr stderr . T.pack
newline :: IO ()
newline = IO.putStrLn T.empty
putEntries :: Int -> OEISEntries -> IO ()
putEntries ncols = mapM_ $ \line -> do
case words line of
[] -> return ()
[w] -> put w -- Should never be reached
(key:aNum:rest) -> do
setSGR [ SetColor Foreground Dull Green ]
put key
setSGR [ SetColor Foreground Dull Yellow ]
put $ ' ' : aNum
setSGR []
let crop = if key == "S" then cropSeq else cropLine
put $ ' ' : crop ncols (unwords rest)
newline
updateCache :: FilePath -> IO ()
updateCache home = do
createDirectoryIfMissing False dir
put msgDownloadingCache
get cacheURL >>= BL.writeFile (dir </> cacheFile)
put msgCacheIsUpToDate
where
dir = home </> cacheDir
readCache :: FilePath -> IO Text
readCache home = do
updated <- doesFileExist fname
if updated
then do
age <- liftM2 diffUTCTime getCurrentTime (getModificationTime fname)
when (age > 100*day) $ putErr msgOldCache
(dropPreamble . decompress) `fmap` BL.readFile fname
else
error msgNoCache
where
day = 60*60*24
fname = home </> cacheDir </> cacheFile
decompress = decodeUtf8 . B.concat . BL.toChunks . GZip.decompress
dropPreamble = T.unlines . drop 4 . T.lines
seqs :: Text -> [Text]
seqs = filter (not . T.null) . map mkSeq . T.lines
where
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-")
filterSeqs :: Bool -> FilePath -> IO ()
filterSeqs invert home = do
cache <- readCache home
let f q = (if invert then not else id) (q `T.isInfixOf` cache)
IO.getContents >>= mapM_ IO.putStrLn . filter f . seqs
args :: Parser Args
args = Args
<$> switch (short 'a' <> long "all" <> help "Print all fields")
<*> strOption
( short 'k'
<> metavar "KEYS"
<> value "SN"
<> help "Keys of fields to print [default: SN]" )
<*> option
( short 'n'
<> metavar "N"
<> value 5
<> help "Fetch at most this many entries [default: 5]" )
<*> switch
( long "invert"
<> help "When used as a filter, return sequences NOT in OEIS" )
<*> switch (long "update" <> help "Update the local sequence cache")
<*> switch (long "url" <> help "Print URLs of found entries")
<*> switch (hidden <> long "version")
<*> many (argument str (metavar "TERMS..."))
sloane :: Args -> IO ()
sloane (Args _ _ _ _ _ _ True _ ) = put name >> newline
sloane (Args _ _ _ _ True _ _ _ ) = getHomeDirectory >>= updateCache
sloane (Args _ _ _ inv _ _ _ []) = getHomeDirectory >>= filterSeqs inv
sloane (Args a keys n _ _ url _ ts) = do
ncols <- getWidth
hits <- searchOEIS n (unwords ts)
let pick = if a then id else select keys
unless (null hits) $ do
newline
if url
then put (urls hits)
else putEntries (ncols - 10) (pick hits)
newline
main :: IO ()
main = execParser (info (h <*> args) (fullDesc <> header name)) >>= sloane
where
h = abortOption ShowHelpText $ hidden <> long "help"