AST.hs

{-
 - Copyright 2024 Tycho Andersen
 -
 - This module exists as a small AST to replace language-python, which seems
 - largely unmaintained at this point. This is not really intended to be a
 - complete python grammer, mostly it is just enough to be what xcffib needs to
 - generate its trees.
 -
 - It has no annotations as in language-python, because xcffib doesn't use them.
 -
 - The grammar not complete: it cannot express anything other than empty
 - dictionaries, cannot express sets at all, cannot do async, etc. all because
 - these features of the language are unused by xcffib.
 -
 - The grammar is not sound: it has one "op" class, representing both unary and
 - binary operators.
 -
 - Use at your own risk :)
 -}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module Data.XCB.Python.AST (
  Expr(..),
  Suite,
  Statement(..),
  Ident,
  Op(..),
  Pretty(..),
  prettyText,
  PseudoExpr(..)
) where

import Prelude hiding ((<>))

import Data.Maybe

import Text.PrettyPrint

type Suite = [Statement]

type Ident = String

-- we don't use that many operations :)
data Op
   = Plus
   | Minus
   | Multiply
   | FloorDivide
   | BinaryAnd
   | ShiftRight
   | ShiftLeft
   | Invert
   | Equality
   | LessThan
   | Modulo
   deriving (Eq, Ord, Show)

data Statement
    = Import
      { import_item :: Ident }
    -- this is only a one level dotted import; other levels are not supported
    | FromImport
      { from_module :: Ident
      , from_item :: Ident
      }
    -- we skip While, For, AsyncFor, etc. because we don't use them
    | Fun
      { fun_name :: Ident
      -- only named variables, no default values, kwargs, etc.
      , fun_args :: [Ident]
      , fun_body :: Suite
      }
    | Decorated
      { decorations :: Ident
      -- only decorated functions/methods are supported
      , fun_name :: Ident
      , fun_args :: [Ident]
      , fun_body :: Suite
      }
    -- skip AsyncFun...
    | Class
      { class_name :: Ident
      -- same as functions, only identifiers allowed
      , class_args :: [Ident]
      , class_body :: Suite
      }
    | Conditional
      { if_cond :: Expr
      , if_body :: Suite
      , else_body :: Suite
      }
    | Assign
      { assign_to :: Expr
      , assign_expr :: Expr
      }
    | AugmentedAssign
      { aug_assign_to :: Expr
      , aug_assign_op :: Op
      , aug_assign_expr :: Expr
      }
    -- skip AnnotatedAssign, Decorated
    | Return
      { return_expr :: Maybe (Expr) }
    -- skip Try, With, AsyncWith
    | Pass {}
    -- skip Break, Continue, Delete
    | StmtExpr
      { stmt_expr :: Expr }
    -- skip Global, NonLocal, Assert, Print, Exec
    | Raise { raise_exception :: Ident }
    deriving (Eq, Ord, Show)

data Expr
    = Var { var :: Ident }
    | Int { int_value :: Int }
    | Bool { bool_value :: Bool }
    | None
    | Strings { strings_strings :: [String] }
    | Call
      { call_fun :: Expr
      , call_args :: [Expr]
      }
    | CondExpr
      { ce_true_branch :: Expr
      , ce_conditon :: Expr
      , ce_false_branch :: Expr
      }
    | Subscript
      { subscriptee :: Expr
      , subscript_expr :: Expr
      }
    | BinaryOp
      { operator :: Op
      , binop_left :: Expr
      , binop_right :: Expr
      }
    | UnaryOp
      { operator :: Op
      , unop_arg :: Expr
      }
    | Dot
      { dot_expr :: Expr
      , dot_attribute :: Ident
      }
    | Tuple { tuple_exprs :: [Expr] }
    | EmptyDict {} -- an empty dictionary
    | Paren { paren_expr :: Expr }
    deriving (Eq, Ord, Show)

prettyText :: Pretty a => a -> String
prettyText = render . pretty

class Pretty a where
    pretty :: a -> Doc

instance Pretty Op where
    pretty Plus = text "+"
    pretty Minus = text "-"
    pretty Multiply = text "*"
    pretty FloorDivide = text "//"
    pretty BinaryAnd = text "&"
    pretty ShiftRight = text ">>"
    pretty ShiftLeft = text "<<"
    pretty Invert = text "~"
    pretty Equality = text "=="
    pretty LessThan = text "<"
    pretty Modulo = text "%"

_reserved :: [String]
_reserved = [ "None"
            , "def"
            , "class"
            , "and"
            , "or"
            ]

-- | Sanitize identifiers.
instance Pretty Ident where
    pretty s | s `elem` _reserved = text $ "_" ++ s
    pretty s | isInt s = text $ "_" ++ s
      where
        isInt str = isJust $ ((maybeRead str) :: Maybe Int)
        maybeRead = fmap fst . listToMaybe . reads
    pretty s = text s

instance Pretty Suite where
    pretty stmts = vcat $ map pretty stmts

indent :: Doc -> Doc
indent = nest 4

instance Pretty Statement where
    pretty (Import item) = text "import" <+> pretty item
    pretty (FromImport source item) = text "from" <+> text source <+> text "import" <+> pretty item
    pretty (Fun name args bod) = text "def" <+> pretty name <> (parens (addCommas args)) <> colon
                                 $+$ indent (pretty bod)
    pretty (Decorated decorator name args bod) = text "@" <> text decorator $+$ pretty (Fun name args bod)
    pretty (Class name [] body) = text "@dataclass(init=False)" $+$ text "class" <+> pretty name <> colon $+$ indent (pretty body)
    pretty (Class name superclasses body) = text "@dataclass(init=False)" $+$ text "class" <+> pretty name <> parens (addCommas superclasses) <> colon $+$ indent (pretty body)
    pretty (Conditional cond if_ else_) = text "if" <+> pretty cond <> colon $+$ indent (pretty if_) $+$ pretty else_
    pretty (Assign to expr) = pretty to <+> text "=" <+> pretty expr
    pretty (AugmentedAssign to op expr) = pretty to <+> pretty op <> text "=" <+> pretty expr
    pretty (Return (Just expr)) = text "return" <+> pretty expr
    pretty (Return Nothing) = text "return"
    pretty Pass = text "pass"
    pretty (StmtExpr expr) = pretty expr
    pretty (Raise exc) = text "raise" <+> pretty exc

class PseudoExpr a where
  getExpr :: a -> Expr

instance PseudoExpr String where
  getExpr s = Var s
instance PseudoExpr Expr where
  getExpr = id

addCommas :: PseudoExpr a => [a] -> Doc
addCommas exprs = hsep $ punctuate (text ",") (map (pretty . getExpr) exprs)

instance Pretty Expr where
    pretty (Var v) = pretty v
    pretty (Int i) = integer (toInteger i)
    pretty (Bool True) = text "True"
    pretty (Bool False) = text "False"
    pretty None = text "None"
    pretty (Strings xs) = hcat (map text xs)
    pretty (Call fun args) = pretty fun <> lparen <> addCommas args <> rparen
    pretty (CondExpr trueB cond falseB) = pretty trueB <+> text "if" <+> pretty cond <+> text "else" <+> pretty falseB
    pretty (Subscript thing expr) = pretty thing <> brackets (pretty expr)
    pretty (BinaryOp op left right) = pretty left <+> pretty op <+> pretty right
    pretty (UnaryOp op arg) = pretty op <> pretty arg
    pretty (Dot thing attr) = pretty thing <> text "." <> pretty attr
    pretty (Tuple [expr]) = pretty expr <> comma -- one element still needs to be a tuple
    pretty (Tuple exprs) = addCommas exprs
    pretty EmptyDict = text "{}"
    pretty (Paren expr) = lparen <> pretty expr <> rparen